-
Notifications
You must be signed in to change notification settings - Fork 1
/
manager.py
54 lines (46 loc) · 1.93 KB
/
manager.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import numpy as np
import supervision as sv
from dataclasses import dataclass, field
from typing import List, Dict, Set
from uuid import uuid4, UUID
@dataclass
class Polygon:
id: UUID
color: sv.Color
is_output: bool
zone: sv.PolygonZone
detected_ids: set[int] = field(default_factory=set)
class ZonesManager:
def __init__(self):
self.input_zones: List[Polygon] = list()
self.output_zones: List[Polygon] = list()
self.tracker: Dict[UUID, Dict[UUID, Set[int]]] = dict()
def create_polygon(self, coords: np.ndarray, zone_is_output: bool, video_info: sv.VideoInfo) -> Polygon:
zone = sv.PolygonZone(
coords,
(video_info.width, video_info.height),
sv.Position.CENTER
)
zone_array = self.output_zones if zone_is_output else self.input_zones
polygon = Polygon(
id=uuid4(),
color=sv.Color(np.random.choice(256), np.random.choice(256), np.random.choice(256)),
is_output=zone_is_output,
zone=zone
)
zone_array.append(polygon)
if polygon.is_output:
self.tracker.setdefault(polygon.id, {})
return polygon
def update(self, input_zones: List[sv.Detections], output_zones: List[sv.Detections]):
for polygon in self.input_zones:
for tracker_id in input_zones[self.input_zones.index(polygon)].tracker_id:
polygon.detected_ids.add(tracker_id)
for polygon in self.output_zones:
for tracker_id in output_zones[self.output_zones.index(polygon)].tracker_id:
polygon.detected_ids.add(tracker_id)
for z in self.input_zones:
if tracker_id in z.detected_ids:
self.tracker.setdefault(polygon.id, {})
self.tracker[polygon.id].setdefault(z.id, set())
self.tracker[polygon.id][z.id].add(tracker_id)