Browse Source

safe refactoring (#2552)

Co-authored-by: YS <ys@gm.com>
Yuriy Sannikov 3 years ago
parent
commit
80627e4989
7 changed files with 36 additions and 36 deletions
  1. 2 3
      frigate/config.py
  2. 1 21
      frigate/edgetpu.py
  3. 3 2
      frigate/http.py
  4. 1 1
      frigate/mqtt.py
  5. 3 2
      frigate/object_processing.py
  6. 22 4
      frigate/util.py
  7. 4 3
      frigate/video.py

+ 2 - 3
frigate/config.py

@@ -13,8 +13,7 @@ from pydantic import BaseModel, Extra, Field, validator
 from pydantic.fields import PrivateAttr
 from pydantic.fields import PrivateAttr
 
 
 from frigate.const import BASE_DIR, CACHE_DIR, YAML_EXT
 from frigate.const import BASE_DIR, CACHE_DIR, YAML_EXT
-from frigate.edgetpu import load_labels
-from frigate.util import create_mask, deep_merge
+from frigate.util import create_mask, deep_merge, load_labels
 
 
 logger = logging.getLogger(__name__)
 logger = logging.getLogger(__name__)
 
 
@@ -640,7 +639,7 @@ class ModelConfig(FrigateBaseModel):
         return self._merged_labelmap
         return self._merged_labelmap
 
 
     @property
     @property
-    def colormap(self) -> Dict[int, tuple[int, int, int]]:
+    def colormap(self) -> Dict[int, Tuple[int, int, int]]:
         return self._colormap
         return self._colormap
 
 
     def __init__(self, **config):
     def __init__(self, **config):

+ 1 - 21
frigate/edgetpu.py

@@ -13,31 +13,11 @@ import tflite_runtime.interpreter as tflite
 from setproctitle import setproctitle
 from setproctitle import setproctitle
 from tflite_runtime.interpreter import load_delegate
 from tflite_runtime.interpreter import load_delegate
 
 
-from frigate.util import EventsPerSecond, SharedMemoryFrameManager, listen
+from frigate.util import EventsPerSecond, SharedMemoryFrameManager, listen, load_labels
 
 
 logger = logging.getLogger(__name__)
 logger = logging.getLogger(__name__)
 
 
 
 
-def load_labels(path, encoding="utf-8"):
-    """Loads labels from file (with or without index numbers).
-    Args:
-      path: path to label file.
-      encoding: label file encoding.
-    Returns:
-      Dictionary mapping indices to labels.
-    """
-    with open(path, "r", encoding=encoding) as f:
-        lines = f.readlines()
-        if not lines:
-            return {}
-
-        if lines[0].split(" ", maxsplit=1)[0].isdigit():
-            pairs = [line.split(" ", maxsplit=1) for line in lines]
-            return {int(index): label.strip() for index, label in pairs}
-        else:
-            return {index: line.strip() for index, line in enumerate(lines)}
-
-
 class ObjectDetector(ABC):
 class ObjectDetector(ABC):
     @abstractmethod
     @abstractmethod
     def detect(self, tensor_input, threshold=0.4):
     def detect(self, tensor_input, threshold=0.4):

+ 3 - 2
frigate/http.py

@@ -359,9 +359,10 @@ def best(camera_name, label):
 
 
         crop = bool(request.args.get("crop", 0, type=int))
         crop = bool(request.args.get("crop", 0, type=int))
         if crop:
         if crop:
-            box = best_object.get("box", (0, 0, 300, 300))
+            box_size = 300
+            box = best_object.get("box", (0, 0, box_size, box_size))
             region = calculate_region(
             region = calculate_region(
-                best_frame.shape, box[0], box[1], box[2], box[3], 1.1
+                best_frame.shape, box[0], box[1], box[2], box[3], box_size, multiplier=1.1
             )
             )
             best_frame = best_frame[region[1] : region[3], region[0] : region[2]]
             best_frame = best_frame[region[1] : region[3], region[0] : region[2]]
 
 

+ 1 - 1
frigate/mqtt.py

@@ -107,7 +107,7 @@ def create_mqtt_client(config: FrigateConfig, camera_metrics):
                     + str(rc)
                     + str(rc)
                 )
                 )
 
 
-        logger.info("MQTT connected")
+        logger.debug("MQTT connected")
         client.subscribe(f"{mqtt_config.topic_prefix}/#")
         client.subscribe(f"{mqtt_config.topic_prefix}/#")
         client.publish(mqtt_config.topic_prefix + "/available", "online", retain=True)
         client.publish(mqtt_config.topic_prefix + "/available", "online", retain=True)
 
 

+ 3 - 2
frigate/object_processing.py

@@ -18,12 +18,12 @@ import numpy as np
 
 
 from frigate.config import CameraConfig, SnapshotsConfig, RecordConfig, FrigateConfig
 from frigate.config import CameraConfig, SnapshotsConfig, RecordConfig, FrigateConfig
 from frigate.const import CACHE_DIR, CLIPS_DIR, RECORD_DIR
 from frigate.const import CACHE_DIR, CLIPS_DIR, RECORD_DIR
-from frigate.edgetpu import load_labels
 from frigate.util import (
 from frigate.util import (
     SharedMemoryFrameManager,
     SharedMemoryFrameManager,
     calculate_region,
     calculate_region,
     draw_box_with_label,
     draw_box_with_label,
     draw_timestamp,
     draw_timestamp,
+    load_labels,
 )
 )
 
 
 logger = logging.getLogger(__name__)
 logger = logging.getLogger(__name__)
@@ -264,8 +264,9 @@ class TrackedObject:
 
 
         if crop:
         if crop:
             box = self.thumbnail_data["box"]
             box = self.thumbnail_data["box"]
+            box_size = 300
             region = calculate_region(
             region = calculate_region(
-                best_frame.shape, box[0], box[1], box[2], box[3], 1.1
+                best_frame.shape, box[0], box[1], box[2], box[3], box_size, multiplier=1.1
             )
             )
             best_frame = best_frame[region[1] : region[3], region[0] : region[2]]
             best_frame = best_frame[region[1] : region[3], region[0] : region[2]]
 
 

+ 22 - 4
frigate/util.py

@@ -189,12 +189,12 @@ def draw_box_with_label(
     )
     )
 
 
 
 
-def calculate_region(frame_shape, xmin, ymin, xmax, ymax, multiplier=2):
+def calculate_region(frame_shape, xmin, ymin, xmax, ymax, model_size, multiplier=2):
     # size is the longest edge and divisible by 4
     # size is the longest edge and divisible by 4
     size = int((max(xmax - xmin, ymax - ymin) * multiplier) // 4 * 4)
     size = int((max(xmax - xmin, ymax - ymin) * multiplier) // 4 * 4)
-    # dont go any smaller than 300
-    if size < 300:
-        size = 300
+    # dont go any smaller than the model_size
+    if size < model_size:
+        size = model_size
 
 
     # x_offset is midpoint of bounding box minus half the size
     # x_offset is midpoint of bounding box minus half the size
     x_offset = int((xmax - xmin) / 2.0 + xmin - size / 2.0)
     x_offset = int((xmax - xmin) / 2.0 + xmin - size / 2.0)
@@ -601,6 +601,24 @@ def add_mask(mask, mask_img):
     )
     )
     cv2.fillPoly(mask_img, pts=[contour], color=(0))
     cv2.fillPoly(mask_img, pts=[contour], color=(0))
 
 
+def load_labels(path, encoding="utf-8"):
+    """Loads labels from file (with or without index numbers).
+    Args:
+      path: path to label file.
+      encoding: label file encoding.
+    Returns:
+      Dictionary mapping indices to labels.
+    """
+    with open(path, "r", encoding=encoding) as f:
+        lines = f.readlines()
+        if not lines:
+            return {}
+
+        if lines[0].split(" ", maxsplit=1)[0].isdigit():
+            pairs = [line.split(" ", maxsplit=1) for line in lines]
+            return {int(index): label.strip() for index, label in pairs}
+        else:
+            return {index: line.strip() for index, line in enumerate(lines)}
 
 
 class FrameManager(ABC):
 class FrameManager(ABC):
     @abstractmethod
     @abstractmethod

+ 4 - 3
frigate/video.py

@@ -529,15 +529,16 @@ def process_frames(
         # combine motion boxes with known locations of existing objects
         # combine motion boxes with known locations of existing objects
         combined_boxes = reduce_boxes(motion_boxes + tracked_object_boxes)
         combined_boxes = reduce_boxes(motion_boxes + tracked_object_boxes)
 
 
+        region_min_size = max(model_shape[0], model_shape[1])
         # compute regions
         # compute regions
         regions = [
         regions = [
-            calculate_region(frame_shape, a[0], a[1], a[2], a[3], 1.2)
+            calculate_region(frame_shape, a[0], a[1], a[2], a[3], region_min_size, multiplier=1.2)
             for a in combined_boxes
             for a in combined_boxes
         ]
         ]
 
 
         # consolidate regions with heavy overlap
         # consolidate regions with heavy overlap
         regions = [
         regions = [
-            calculate_region(frame_shape, a[0], a[1], a[2], a[3], 1.0)
+            calculate_region(frame_shape, a[0], a[1], a[2], a[3], region_min_size, multiplier=1.0)
             for a in reduce_boxes(regions, 0.4)
             for a in reduce_boxes(regions, 0.4)
         ]
         ]
 
 
@@ -596,7 +597,7 @@ def process_frames(
                         box = obj[2]
                         box = obj[2]
                         # calculate a new region that will hopefully get the entire object
                         # calculate a new region that will hopefully get the entire object
                         region = calculate_region(
                         region = calculate_region(
-                            frame_shape, box[0], box[1], box[2], box[3]
+                            frame_shape, box[0], box[1], box[2], box[3], region_min_size
                         )
                         )
 
 
                         regions.append(region)
                         regions.append(region)