WIP: convert to camera class

frigate/video.py

@@ -1,8 +1,14 @@
+import os
 import time
 import datetime
 import cv2
 import threading
+import ctypes
+import multiprocessing as mp
 from . util import tonumpyarray
+from . object_detection import FramePrepper
+from . objects import ObjectCleaner, ObjectParser, BestPersonFrame
+from . mqtt import MqttObjectPublisher
 
 # fetch the frames as fast a possible, only decoding the frames when the
 # detection_process has consumed the current frame
@@ -85,3 +91,130 @@ class FrameTracker(threading.Thread):
             for k in stored_frame_times:
                 if (now - k) > 2:
                     del self.recent_frames[k]
+
+def get_frame_shape(rtsp_url):
+    # capture a single frame and check the frame shape so the correct array
+    # size can be allocated in memory
+    video = cv2.VideoCapture(rtsp_url)
+    ret, frame = video.read()
+    frame_shape = frame.shape
+    video.release()
+    return frame_shape
+
+def get_rtsp_url(rtsp_config):
+    if (rtsp_config['password'].startswith('$')):
+        rtsp_config['password'] = os.getenv(rtsp_config['password'][1:])
+    return 'rtsp://{}:{}@{}:{}{}'.format(rtsp_config['user'], 
+        rtsp_config['password'], rtsp_config['host'], rtsp_config['port'],
+        rtsp_config['path'])
+
+class Camera:
+    def __init__(self, name, config, prepped_frame_queue, mqtt_client, mqtt_prefix):
+        self.name = name
+        self.config = config
+        self.detected_objects = []
+        self.recent_frames = {}
+        self.rtsp_url = get_rtsp_url(self.config['rtsp'])
+        self.regions = self.config['regions']
+        self.frame_shape = get_frame_shape(self.rtsp_url)
+        self.mqtt_client = mqtt_client
+        self.mqtt_topic_prefix = '{}/{}'.format(mqtt_prefix, self.name)
+
+        # compute the flattened array length from the shape of the frame
+        flat_array_length = self.frame_shape[0] * self.frame_shape[1] * self.frame_shape[2]
+        # create shared array for storing the full frame image data
+        self.shared_frame_array = mp.Array(ctypes.c_uint8, flat_array_length)
+        # create shared value for storing the frame_time
+        self.shared_frame_time = mp.Value('d', 0.0)
+        # Lock to control access to the frame
+        self.frame_lock = mp.Lock()
+        # Condition for notifying that a new frame is ready
+        self.frame_ready = mp.Condition()
+        # Condition for notifying that objects were parsed
+        self.objects_parsed = mp.Condition()
+
+        # shape current frame so it can be treated as a numpy image
+        self.shared_frame_np = tonumpyarray(self.shared_frame_array).reshape(self.frame_shape)
+
+        # create the process to capture frames from the RTSP stream and store in a shared array
+        self.capture_process = mp.Process(target=fetch_frames, args=(self.shared_frame_array, 
+            self.shared_frame_time, self.frame_lock, self.frame_ready, self.frame_shape, self.rtsp_url))
+        self.capture_process.daemon = True
+
+        # for each region, create a separate thread to resize the region and prep for detection
+        self.detection_prep_threads = []
+        for region in self.config['regions']:
+            self.detection_prep_threads.append(FramePrepper(
+                self.name,
+                self.shared_frame_np,
+                self.shared_frame_time,
+                self.frame_ready,
+                self.frame_lock,
+                region['size'], region['x_offset'], region['y_offset'],
+                prepped_frame_queue
+            ))
+        
+        # start a thread to store recent motion frames for processing
+        self.frame_tracker = FrameTracker(self.shared_frame_np, self.shared_frame_time, 
+            self.frame_ready, self.frame_lock, self.recent_frames)
+        self.frame_tracker.start()
+
+        # start a thread to store the highest scoring recent person frame
+        self.best_person_frame = BestPersonFrame(self.objects_parsed, self.recent_frames, self.detected_objects)
+        self.best_person_frame.start()
+
+        # start a thread to expire objects from the detected objects list
+        self.object_cleaner = ObjectCleaner(self.objects_parsed, self.detected_objects)
+        self.object_cleaner.start()
+
+        # start a thread to publish object scores (currently only person)
+        mqtt_publisher = MqttObjectPublisher(self.mqtt_client, self.mqtt_topic_prefix, self.objects_parsed, self.detected_objects)
+        mqtt_publisher.start()
+    
+    def start(self):
+        self.capture_process.start()
+        # start the object detection prep threads
+        for detection_prep_thread in self.detection_prep_threads:
+            detection_prep_thread.start()
+    
+    def join(self):
+        self.capture_process.join()
+    
+    def get_capture_pid(self):
+        return self.capture_process.pid
+    
+    def add_objects(self, objects):
+        if len(objects) == 0:
+            return
+
+        for obj in objects:
+            if obj['name'] == 'person':
+                person_area = (obj['xmax']-obj['xmin'])*(obj['ymax']-obj['ymin'])
+                # find the matching region
+                region = None
+                for r in self.regions:
+                    if (
+                            obj['xmin'] >= r['x_offset'] and
+                            obj['ymin'] >= r['y_offset'] and
+                            obj['xmax'] <= r['x_offset']+r['size'] and
+                            obj['ymax'] <= r['y_offset']+r['size']
+                        ): 
+                        region = r
+                        break
+                
+                # if the min person area is larger than the
+                # detected person, don't add it to detected objects
+                if region and region['min_person_area'] > person_area:
+                    continue
+
+            self.detected_objects.append(obj)
+
+        with self.objects_parsed:
+            self.objects_parsed.notify_all()
+
+    def get_best_person(self):
+        return self.best_person_frame.best_frame
+
+
+    
+