forked from blakeblackshear/frigate
-
Notifications
You must be signed in to change notification settings - Fork 0
/
process_clip.py
148 lines (128 loc) · 6.25 KB
/
process_clip.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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
import sys
import click
import os
import datetime
from unittest import TestCase, main
from frigate.video import process_frames, start_or_restart_ffmpeg, capture_frames, get_frame_shape
from frigate.util import DictFrameManager, EventsPerSecond, draw_box_with_label
from frigate.motion import MotionDetector
from frigate.edgetpu import LocalObjectDetector
from frigate.objects import ObjectTracker
import multiprocessing as mp
import numpy as np
import cv2
from frigate.object_processing import COLOR_MAP, CameraState
class ProcessClip():
def __init__(self, clip_path, frame_shape, config):
self.clip_path = clip_path
self.frame_shape = frame_shape
self.camera_name = 'camera'
self.frame_manager = DictFrameManager()
self.frame_queue = mp.Queue()
self.detected_objects_queue = mp.Queue()
self.camera_state = CameraState(self.camera_name, config, self.frame_manager)
def load_frames(self):
fps = EventsPerSecond()
skipped_fps = EventsPerSecond()
stop_event = mp.Event()
detection_frame = mp.Value('d', datetime.datetime.now().timestamp()+100000)
current_frame = mp.Value('d', 0.0)
ffmpeg_cmd = f"ffmpeg -hide_banner -loglevel panic -i {self.clip_path} -f rawvideo -pix_fmt rgb24 pipe:".split(" ")
ffmpeg_process = start_or_restart_ffmpeg(ffmpeg_cmd, self.frame_shape[0]*self.frame_shape[1]*self.frame_shape[2])
capture_frames(ffmpeg_process, self.camera_name, self.frame_shape, self.frame_manager, self.frame_queue, 1, fps, skipped_fps, stop_event, detection_frame, current_frame)
ffmpeg_process.wait()
ffmpeg_process.communicate()
def process_frames(self, objects_to_track=['person'], object_filters={}):
mask = np.zeros((self.frame_shape[0], self.frame_shape[1], 1), np.uint8)
mask[:] = 255
motion_detector = MotionDetector(self.frame_shape, mask)
object_detector = LocalObjectDetector(labels='/labelmap.txt')
object_tracker = ObjectTracker(10)
process_fps = mp.Value('d', 0.0)
detection_fps = mp.Value('d', 0.0)
current_frame = mp.Value('d', 0.0)
stop_event = mp.Event()
process_frames(self.camera_name, self.frame_queue, self.frame_shape, self.frame_manager, motion_detector, object_detector, object_tracker, self.detected_objects_queue,
process_fps, detection_fps, current_frame, objects_to_track, object_filters, mask, stop_event, exit_on_empty=True)
def objects_found(self, debug_path=None):
obj_detected = False
top_computed_score = 0.0
def handle_event(name, obj):
nonlocal obj_detected
nonlocal top_computed_score
if obj['computed_score'] > top_computed_score:
top_computed_score = obj['computed_score']
if not obj['false_positive']:
obj_detected = True
self.camera_state.on('new', handle_event)
self.camera_state.on('update', handle_event)
while(not self.detected_objects_queue.empty()):
camera_name, frame_time, current_tracked_objects = self.detected_objects_queue.get()
if not debug_path is None:
self.save_debug_frame(debug_path, frame_time, current_tracked_objects.values())
self.camera_state.update(frame_time, current_tracked_objects)
for obj in self.camera_state.tracked_objects.values():
print(f"{frame_time}: {obj['id']} - {obj['computed_score']} - {obj['score_history']}")
return {
'object_detected': obj_detected,
'top_score': top_computed_score
}
def save_debug_frame(self, debug_path, frame_time, tracked_objects):
current_frame = self.frame_manager.get(f"{self.camera_name}{frame_time}")
# draw the bounding boxes on the frame
for obj in tracked_objects:
thickness = 2
color = (0,0,175)
if obj['frame_time'] != frame_time:
thickness = 1
color = (255,0,0)
else:
color = (255,255,0)
# draw the bounding boxes on the frame
box = obj['box']
draw_box_with_label(current_frame, box[0], box[1], box[2], box[3], obj['label'], f"{int(obj['score']*100)}% {int(obj['area'])}", thickness=thickness, color=color)
# draw the regions on the frame
region = obj['region']
draw_box_with_label(current_frame, region[0], region[1], region[2], region[3], 'region', "", thickness=1, color=(0,255,0))
cv2.imwrite(f"{os.path.join(debug_path, os.path.basename(self.clip_path))}.{int(frame_time*1000000)}.jpg", cv2.cvtColor(current_frame, cv2.COLOR_RGB2BGR))
@click.command()
@click.option("-p", "--path", required=True, help="Path to clip or directory to test.")
@click.option("-l", "--label", default='person', help="Label name to detect.")
@click.option("-t", "--threshold", default=0.85, help="Threshold value for objects.")
@click.option("--debug-path", default=None, help="Path to output frames for debugging.")
def process(path, label, threshold, debug_path):
clips = []
if os.path.isdir(path):
files = os.listdir(path)
files.sort()
clips = [os.path.join(path, file) for file in files]
elif os.path.isfile(path):
clips.append(path)
config = {
'snapshots': {
'show_timestamp': False,
'draw_zones': False
},
'zones': {},
'objects': {
'track': [label],
'filters': {
'person': {
'threshold': threshold
}
}
}
}
results = []
for c in clips:
frame_shape = get_frame_shape(c)
process_clip = ProcessClip(c, frame_shape, config)
process_clip.load_frames()
process_clip.process_frames(objects_to_track=config['objects']['track'])
results.append((c, process_clip.objects_found(debug_path)))
for result in results:
print(f"{result[0]}: {result[1]}")
positive_count = sum(1 for result in results if result[1]['object_detected'])
print(f"Objects were detected in {positive_count}/{len(results)}({positive_count/len(results)*100:.2f}%) clip(s).")
if __name__ == '__main__':
process()