allow defining model shape and switch to mobiledet as default model

docker/Dockerfile.base

@@ -36,10 +36,9 @@ RUN pip3 install \
 COPY nginx/nginx.conf /etc/nginx/nginx.conf
 
 # get model and labels
-ARG MODEL_REFS=7064b94dd5b996189242320359dbab8b52c94a84
 COPY labelmap.txt /labelmap.txt
-RUN wget -q https://github.com/google-coral/edgetpu/raw/$MODEL_REFS/test_data/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite -O /edgetpu_model.tflite
-RUN wget -q https://github.com/google-coral/edgetpu/raw/$MODEL_REFS/test_data/ssd_mobilenet_v2_coco_quant_postprocess.tflite -O /cpu_model.tflite
+RUN wget -q https://github.com/google-coral/test_data/raw/master/ssdlite_mobiledet_coco_qat_postprocess_edgetpu.tflite -O /edgetpu_model.tflite
+RUN wget -q https://github.com/google-coral/test_data/raw/master/ssdlite_mobiledet_coco_qat_postprocess.tflite -O /cpu_model.tflite
 
 WORKDIR /opt/frigate/
 ADD frigate frigate/

frigate/app.py

@@ -115,27 +115,29 @@ def init_mqtt(self):
  self.mqtt_client = create_mqtt_client(self.config.mqtt)
 
  def start_detectors(self):
+ model_shape = (self.config.model.height, self.config.model.width)
  for name in self.config.cameras.keys():
  self.detection_out_events[name] = mp.Event()
- shm_in = mp.shared_memory.SharedMemory(name=name, create=True, size=300*300*3)
+ shm_in = mp.shared_memory.SharedMemory(name=name, create=True, size=self.config.model.height*self.config.model.width*3)
  shm_out = mp.shared_memory.SharedMemory(name=f"out-{name}", create=True, size=20*6*4)
  self.detection_shms.append(shm_in)
  self.detection_shms.append(shm_out)
 
  for name, detector in self.config.detectors.items():
  if detector.type == 'cpu':
- self.detectors[name] = EdgeTPUProcess(name, self.detection_queue, out_events=self.detection_out_events, tf_device='cpu')
+ self.detectors[name] = EdgeTPUProcess(name, self.detection_queue, self.detection_out_events, model_shape, tf_device='cpu')
  if detector.type == 'edgetpu':
- self.detectors[name] = EdgeTPUProcess(name, self.detection_queue, out_events=self.detection_out_events, tf_device=detector.device)
+ self.detectors[name] = EdgeTPUProcess(name, self.detection_queue, self.detection_out_events, model_shape, tf_device=detector.device)
 
  def start_detected_frames_processor(self):
  self.detected_frames_processor = TrackedObjectProcessor(self.config, self.mqtt_client, self.config.mqtt.topic_prefix, 
  self.detected_frames_queue, self.event_queue, self.event_processed_queue, self.stop_event)
  self.detected_frames_processor.start()
 
  def start_camera_processors(self):
+ model_shape = (self.config.model.height, self.config.model.width)
  for name, config in self.config.cameras.items():
- camera_process = mp.Process(target=track_camera, name=f"camera_processor:{name}", args=(name, config,
+ camera_process = mp.Process(target=track_camera, name=f"camera_processor:{name}", args=(name, config, model_shape,
  self.detection_queue, self.detection_out_events[name], self.detected_frames_queue, 
  self.camera_metrics[name]))
  camera_process.daemon = True

frigate/config.py

@@ -193,7 +193,7 @@ def ensure_zones_and_cameras_have_different_names(cameras):
 
 FRIGATE_CONFIG_SCHEMA = vol.Schema(
  {
- vol.Optional('model', default={'width': 300, 'height': 300}): {
+ vol.Optional('model', default={'width': 320, 'height': 320}): {
  vol.Required('width'): int,
  vol.Required('height'): int
  },

frigate/edgetpu.py

@@ -106,7 +106,7 @@ def detect_raw(self, tensor_input):
 
  return detections
 
-def run_detector(name: str, detection_queue: mp.Queue, out_events: Dict[str, mp.Event], avg_speed, start, tf_device):
+def run_detector(name: str, detection_queue: mp.Queue, out_events: Dict[str, mp.Event], avg_speed, start, model_shape, tf_device):
  threading.current_thread().name = f"detector:{name}"
  logger = logging.getLogger(f"detector.{name}")
  logger.info(f"Starting detection process: {os.getpid()}")
@@ -139,7 +139,7 @@ def receiveSignal(signalNumber, frame):
  connection_id = detection_queue.get(timeout=5)
  except queue.Empty:
  continue
- input_frame = frame_manager.get(connection_id, (1,300,300,3))
+ input_frame = frame_manager.get(connection_id, (1,model_shape[0],model_shape[1],3))
 
  if input_frame is None:
  continue
@@ -155,13 +155,14 @@ def receiveSignal(signalNumber, frame):
  avg_speed.value = (avg_speed.value*9 + duration)/10
 
 class EdgeTPUProcess():
- def __init__(self, name, detection_queue, out_events, tf_device=None):
+ def __init__(self, name, detection_queue, out_events, model_shape, tf_device=None):
  self.name = name
  self.out_events = out_events
  self.detection_queue = detection_queue
  self.avg_inference_speed = mp.Value('d', 0.01)
  self.detection_start = mp.Value('d', 0.0)
  self.detect_process = None
+ self.model_shape = model_shape
  self.tf_device = tf_device
  self.start_or_restart()
 
@@ -178,19 +179,19 @@ def start_or_restart(self):
  self.detection_start.value = 0.0
  if (not self.detect_process is None) and self.detect_process.is_alive():
  self.stop()
- self.detect_process = mp.Process(target=run_detector, name=f"detector:{self.name}", args=(self.name, self.detection_queue, self.out_events, self.avg_inference_speed, self.detection_start, self.tf_device))
+ self.detect_process = mp.Process(target=run_detector, name=f"detector:{self.name}", args=(self.name, self.detection_queue, self.out_events, self.avg_inference_speed, self.detection_start, self.model_shape, self.tf_device))
  self.detect_process.daemon = True
  self.detect_process.start()
 
 class RemoteObjectDetector():
- def __init__(self, name, labels, detection_queue, event):
+ def __init__(self, name, labels, detection_queue, event, model_shape):
  self.labels = load_labels(labels)
  self.name = name
  self.fps = EventsPerSecond()
  self.detection_queue = detection_queue
  self.event = event
  self.shm = mp.shared_memory.SharedMemory(name=self.name, create=False)
- self.np_shm = np.ndarray((1,300,300,3), dtype=np.uint8, buffer=self.shm.buf)
+ self.np_shm = np.ndarray((1,model_shape[0],model_shape[1],3), dtype=np.uint8, buffer=self.shm.buf)
  self.out_shm = mp.shared_memory.SharedMemory(name=f"out-{self.name}", create=False)
  self.out_np_shm = np.ndarray((20,6), dtype=np.float32, buffer=self.out_shm.buf)
 

frigate/video.py

@@ -64,14 +64,14 @@ def filtered(obj, objects_to_track, object_filters, mask=None):
 
  return False
 
-def create_tensor_input(frame, region):
+def create_tensor_input(frame, model_shape, region):
  cropped_frame = yuv_region_2_rgb(frame, region)
 
  # Resize to 300x300 if needed
- if cropped_frame.shape != (300, 300, 3):
- cropped_frame = cv2.resize(cropped_frame, dsize=(300, 300), interpolation=cv2.INTER_LINEAR)
+ if cropped_frame.shape != (model_shape[0], model_shape[1], 3):
+ cropped_frame = cv2.resize(cropped_frame, dsize=model_shape, interpolation=cv2.INTER_LINEAR)
 
- # Expand dimensions since the model expects images to have shape: [1, 300, 300, 3]
+ # Expand dimensions since the model expects images to have shape: [1, height, width, 3]
  return np.expand_dims(cropped_frame, axis=0)
 
 def stop_ffmpeg(ffmpeg_process, logger):
@@ -241,7 +241,7 @@ def receiveSignal(signalNumber, frame):
  camera_watchdog.start()
  camera_watchdog.join()
 
-def track_camera(name, config: CameraConfig, detection_queue, result_connection, detected_objects_queue, process_info):
+def track_camera(name, config: CameraConfig, model_shape, detection_queue, result_connection, detected_objects_queue, process_info):
  stop_event = mp.Event()
  def receiveSignal(signalNumber, frame):
  stop_event.set()
@@ -260,13 +260,13 @@ def receiveSignal(signalNumber, frame):
  mask = config.mask
 
  motion_detector = MotionDetector(frame_shape, mask, resize_factor=6)
- object_detector = RemoteObjectDetector(name, '/labelmap.txt', detection_queue, result_connection)
+ object_detector = RemoteObjectDetector(name, '/labelmap.txt', detection_queue, result_connection, model_shape)
 
  object_tracker = ObjectTracker(10)
 
  frame_manager = SharedMemoryFrameManager()
 
- process_frames(name, frame_queue, frame_shape, frame_manager, motion_detector, object_detector,
+ process_frames(name, frame_queue, frame_shape, model_shape, frame_manager, motion_detector, object_detector,
  object_tracker, detected_objects_queue, process_info, objects_to_track, object_filters, mask, stop_event)
 
  logger.info(f"{name}: exiting subprocess")
@@ -277,8 +277,8 @@ def reduce_boxes(boxes):
  reduced_boxes = cv2.groupRectangles([list(b) for b in itertools.chain(boxes, boxes)], 1, 0.2)[0]
  return [tuple(b) for b in reduced_boxes]
 
-def detect(object_detector, frame, region, objects_to_track, object_filters, mask):
- tensor_input = create_tensor_input(frame, region)
+def detect(object_detector, frame, model_shape, region, objects_to_track, object_filters, mask):
+ tensor_input = create_tensor_input(frame, model_shape, region)
 
  detections = []
  region_detections = object_detector.detect(tensor_input)
@@ -300,7 +300,7 @@ def detect(object_detector, frame, region, objects_to_track, object_filters, mas
  detections.append(det)
  return detections
 
-def process_frames(camera_name: str, frame_queue: mp.Queue, frame_shape, 
+def process_frames(camera_name: str, frame_queue: mp.Queue, frame_shape, model_shape,
  frame_manager: FrameManager, motion_detector: MotionDetector, 
  object_detector: RemoteObjectDetector, object_tracker: ObjectTracker,
  detected_objects_queue: mp.Queue, process_info: Dict,
@@ -357,7 +357,7 @@ def process_frames(camera_name: str, frame_queue: mp.Queue, frame_shape,
  # resize regions and detect
  detections = []
  for region in regions:
- detections.extend(detect(object_detector, frame, region, objects_to_track, object_filters, mask))
+ detections.extend(detect(object_detector, frame, model_shape, region, objects_to_track, object_filters, mask))
 
  #########
  # merge objects, check for clipped objects and look again up to 4 times
@@ -390,7 +390,7 @@ def process_frames(camera_name: str, frame_queue: mp.Queue, frame_shape,
  box[0], box[1],
  box[2], box[3])
 
- selected_objects.extend(detect(object_detector, frame, region, objects_to_track, object_filters, mask))
+ selected_objects.extend(detect(object_detector, frame, model_shape, region, objects_to_track, object_filters, mask))
 
  refining = True
  else: