-
Notifications
You must be signed in to change notification settings - Fork 0
/
time_spent.py
71 lines (60 loc) · 2.65 KB
/
time_spent.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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
from ultralytics import YOLO
from typing import Self,Tuple
from ultralytics.engine.results import Results
import numpy
from typing import Dict,List,Tuple,Optional
from ..video_loader import VideoLoader
from dataclasses import dataclass
from PIL import Image, ImageDraw
@dataclass
class BBBox:
box:numpy.ndarray
location:str
def __hash__(self) -> int:
return hash(self.location)
def common_area(bb1:numpy.ndarray,bb2:numpy.ndarray)->Optional[Tuple[float,numpy.ndarray,numpy.ndarray]]:
nc = numpy.concatenate([bb1,bb2],axis=0)
max =numpy.max(nc,axis=0)[:2]
min = numpy.min(nc,axis=0)[2:]
if max[0]>min[0] or max[1]>min[1]:
return None
nc = numpy.concatenate([nc,numpy.concatenate([max,min]).reshape(1,-1)],axis=0)
areas = (nc[:,2]-nc[:,0])*(nc[:,3]-nc[:,1])
proprtion = areas[2]/(areas[0]+areas[1]-areas[2])
return (proprtion,areas,nc)
class TimeSpent:
def __init__(self,model:YOLO,treshold=0.4 ) -> None:
self.model=model
self.treshold=treshold
@classmethod
def from_yolo(cls,version:str)->Self:
yolo = YOLO(version)
return cls(yolo)
def time_spent(self,video:VideoLoader,bounding_box:List[BBBox]):
track_ids:Dict[BBBox,Dict[int,Tuple[List[int],Tuple]]]={}
rev_frames:List[List]=[]
iterator = video.frame_iterator()
for index,frame in enumerate(iterator):
img = Image.fromarray(frame)
drawer = ImageDraw.Draw(img)
humans_list=list()
analyzed_frame:Results = self.model.track(frame,persist=True)[0]
for ref_box in bounding_box:
humans = track_ids.setdefault(ref_box,dict())
for box in analyzed_frame.boxes:
if not (box.cls==0).item():
continue
if not box.id:
continue
human_id=box.id.int().item()
frames,color = humans.setdefault(human_id,(list(),tuple(numpy.random.choice(range(256), size=3))))
areas = common_area(ref_box.box,box.xyxy)
if areas is not None:
if areas[1][2]/areas[1][0]>self.treshold:
drawer.rectangle(areas[2][2].tolist(),fill=color)
drawer.rectangle(areas[2][0].tolist(),outline="red")
drawer.rectangle(areas[2][1].tolist(),outline="blue")
humans_list.append(human_id)
frames.append(index)
rev_frames.append((humans_list,numpy.asarray(img)))
return track_ids,rev_frames