Use cv2 to draw object boxes and labels instead of tensorflow object_detection.utils

This facilitates removing tensorflow models and protobuf-python from the
docker image greatly reducing image build time and image size by ~1.3GB

Dockerfile

@@ -50,21 +50,6 @@ RUN  pip install -U pip \
  paho-mqtt \
  PyYAML
 
-# Install tensorflow models object detection
-RUN GIT_SSL_NO_VERIFY=true git clone -q https://github.com/tensorflow/models /usr/local/lib/python3.5/dist-packages/tensorflow/models
-RUN wget -q -P /usr/local/src/ --no-check-certificate https://github.com/google/protobuf/releases/download/v3.5.1/protobuf-python-3.5.1.tar.gz
-
-# Download & build protobuf-python
-RUN cd /usr/local/src/ \
- && tar xf protobuf-python-3.5.1.tar.gz \
- && rm protobuf-python-3.5.1.tar.gz \
- && cd /usr/local/src/protobuf-3.5.1/ \
- && ./configure \
- && make \
- && make install \
- && ldconfig \
- && rm -rf /usr/local/src/protobuf-3.5.1/
-
 # Download & build OpenCV
 RUN wget -q -P /usr/local/src/ --no-check-certificate https://github.com/opencv/opencv/archive/4.0.1.zip
 RUN cd /usr/local/src/ \
@@ -97,10 +82,6 @@ RUN ln -s /python-tflite-source/edgetpu/test_data/coco_labels.txt /coco_labels.t
 RUN (apt-get autoremove -y; \
      apt-get autoclean -y)
 
-# Set TF object detection available
-ENV PYTHONPATH "$PYTHONPATH:/usr/local/lib/python3.5/dist-packages/tensorflow/models/research:/usr/local/lib/python3.5/dist-packages/tensorflow/models/research/slim"
-RUN cd /usr/local/lib/python3.5/dist-packages/tensorflow/models/research && protoc object_detection/protos/*.proto --python_out=.
-
 WORKDIR /opt/frigate/
 ADD frigate frigate/
 COPY detect_objects.py .

frigate/objects.py

@@ -2,7 +2,7 @@ import time
 import datetime
 import threading
 import cv2
-from object_detection.utils import visualization_utils as vis_util
+from . util import drawobjectbox
 
 class ObjectCleaner(threading.Thread):
     def __init__(self, objects_parsed, detected_objects):
@@ -79,18 +79,15 @@ class BestPersonFrame(threading.Thread):
             recent_frames = self.recent_frames.copy()
 
             if not self.best_person is None and self.best_person['frame_time'] in recent_frames:
-                best_frame = recent_frames[self.best_person['frame_time']]
-                best_frame = cv2.cvtColor(best_frame, cv2.COLOR_BGR2RGB)
-                # draw the bounding box on the frame
-                vis_util.draw_bounding_box_on_image_array(best_frame,
-                    self.best_person['ymin'],
-                    self.best_person['xmin'],
-                    self.best_person['ymax'],
-                    self.best_person['xmax'],
-                    color='red',
-                    thickness=2,
-                    display_str_list=["{}: {}%".format(self.best_person['name'],int(self.best_person['score']*100))],
-                    use_normalized_coordinates=False)
+                self.best_frame = recent_frames[self.best_person['frame_time']]
 
-                # convert back to BGR
-                self.best_frame = cv2.cvtColor(best_frame, cv2.COLOR_RGB2BGR)
+                label = "{}: {}%".format(self.best_person['name'], int(self.best_person['score'] * 100))
+                bounding_box = (
+                    self.best_person['xmin'],
+                    self.best_person['ymin'],
+                    self.best_person['xmax'],
+                    self.best_person['ymax']
+                )
+
+                # draw the bounding box on the frame
+                drawobjectbox(self.best_frame, label, bounding_box)

frigate/util.py

@@ -1,5 +1,25 @@
 import numpy as np
+import cv2
+
+LABEL_FONT = cv2.FONT_HERSHEY_PLAIN
+FONT_SCALE = 0.8
 
 # convert shared memory array into numpy array
 def tonumpyarray(mp_arr):
     return np.frombuffer(mp_arr.get_obj(), dtype=np.uint8)
+
+# draw a box with text in the upper left on the image
+def drawobjectbox(image, text, rect):
+    x1, y1, x2, y2 = rect
+
+    # draw the red bounding box
+    cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 2)
+
+    # get the size of the text
+    (text_width, text_height) = cv2.getTextSize(text, LABEL_FONT, FONT_SCALE, 1)[0]
+
+    # draw the text background with padding
+    cv2.rectangle(image, (x1, y1), (x1 + text_width + 8, y1 - text_height - 8), (0, 0, 255), cv2.FILLED)
+
+    # draw the text
+    cv2.putText(image, text, (x1 + 4, y1 - 4), LABEL_FONT, FONT_SCALE, (0, 0, 0), lineType=cv2.LINE_AA)

frigate/video.py

@@ -6,8 +6,8 @@ import threading
 import ctypes
 import multiprocessing as mp
 import numpy as np
-from object_detection.utils import visualization_utils as vis_util
 from . util import tonumpyarray
+from . util import drawobjectbox
 from . object_detection import FramePrepper
 from . objects import ObjectCleaner, BestPersonFrame
 from . mqtt import MqttObjectPublisher
@@ -279,19 +279,16 @@ class Camera:
         with self.frame_lock:
             frame = self.shared_frame_np.copy()
 
-        # convert to RGB for drawing
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        # draw the bounding boxes on the screen
         for obj in detected_objects:
-            vis_util.draw_bounding_box_on_image_array(frame,
-                obj['ymin'],
+            label = "{}: {}%".format(obj['name'], int(obj['score'] * 100))
+            bounding_box = (
                 obj['xmin'],
-                obj['ymax'],
+                obj['ymin'],
                 obj['xmax'],
-                color='red',
-                thickness=2,
-                display_str_list=["{}: {}%".format(obj['name'],int(obj['score']*100))],
-                use_normalized_coordinates=False)
+                obj['ymax']
+            )
+
+            drawobjectbox(frame, label, bounding_box)
 
         for region in self.regions:
             color = (255,255,255)
@@ -299,11 +296,4 @@ class Camera:
                 (region['x_offset']+region['size'], region['y_offset']+region['size']),
                 color, 2)
 
-        # convert back to BGR
-        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
-
         return frame
-
-
-    
-