diff --git a/frigate/track/__init__.py b/frigate/track/__init__.py new file mode 100644 index 000000000..e50980ad2 --- /dev/null +++ b/frigate/track/__init__.py @@ -0,0 +1,13 @@ +from abc import ABC, abstractmethod + +from frigate.config import DetectConfig + +class ObjectTracker(ABC): + + @abstractmethod + def __init__(self, config: DetectConfig): + pass + + @abstractmethod + def match_and_update(self, detections): + pass \ No newline at end of file diff --git a/frigate/objects.py b/frigate/track/centroid_tracker.py similarity index 94% rename from frigate/objects.py rename to frigate/track/centroid_tracker.py index ed04ae191..c2e327869 100644 --- a/frigate/objects.py +++ b/frigate/track/centroid_tracker.py @@ -1,23 +1,15 @@ -import copy -import datetime -import itertools -import multiprocessing as mp +from collections import defaultdict import random import string -import threading -import time -from collections import defaultdict - -import cv2 import numpy as np from scipy.spatial import distance as dist - from frigate.config import DetectConfig + +from frigate.track import ObjectTracker from frigate.util import intersection_over_union - -class ObjectTracker: - def __init__(self, config: DetectConfig): +class CentroidTracker(ObjectTracker): + def __init__(self, config: DetectConfig): self.tracked_objects = {} self.disappeared = {} self.positions = {} @@ -141,11 +133,11 @@ class ObjectTracker: if self.is_expired(id): self.deregister(id) - def match_and_update(self, frame_time, new_objects): + def match_and_update(self, frame_time, detections): # group by name - new_object_groups = defaultdict(lambda: []) - for obj in new_objects: - new_object_groups[obj[0]].append( + detection_groups = defaultdict(lambda: []) + for obj in detections: + detection_groups[obj[0]].append( { "label": obj[0], "score": obj[1], @@ -160,17 +152,17 @@ class ObjectTracker: # update any tracked objects with labels that are not # seen in the current objects and deregister if needed for obj in list(self.tracked_objects.values()): - if not obj["label"] in new_object_groups: + if not obj["label"] in detection_groups: if self.disappeared[obj["id"]] >= self.max_disappeared: self.deregister(obj["id"]) else: self.disappeared[obj["id"]] += 1 - if len(new_objects) == 0: + if len(detections) == 0: return # track objects for each label type - for label, group in new_object_groups.items(): + for label, group in detection_groups.items(): current_objects = [ o for o in self.tracked_objects.values() if o["label"] == label ] @@ -243,4 +235,4 @@ class ObjectTracker: # register each new input centroid as a trackable object else: for col in unusedCols: - self.register(col, group[col]) + self.register(col, group[col]) \ No newline at end of file diff --git a/frigate/video.py b/frigate/video.py index a3356a586..abdfafba9 100755 --- a/frigate/video.py +++ b/frigate/video.py @@ -19,7 +19,8 @@ from frigate.const import CACHE_DIR from frigate.object_detection import RemoteObjectDetector from frigate.log import LogPipe from frigate.motion import MotionDetector -from frigate.objects import ObjectTracker +from frigate.track import ObjectTracker +from frigate.track.centroid_tracker import CentroidTracker from frigate.util import ( EventsPerSecond, FrameManager, @@ -472,7 +473,7 @@ def track_camera( name, labelmap, detection_queue, result_connection, model_config, stop_event ) - object_tracker = ObjectTracker(config.detect) + object_tracker = CentroidTracker(config.detect) frame_manager = SharedMemoryFrameManager() diff --git a/process_clip.py b/process_clip.py index d8dabbedd..4aa0fa9b3 100644 --- a/process_clip.py +++ b/process_clip.py @@ -1,7 +1,9 @@ import sys from typing_extensions import runtime -sys.path.append("/lab/frigate") +from frigate.track.centroid_tracker import CentroidTracker + +sys.path.append("/workspace/frigate") import json import logging @@ -19,7 +21,6 @@ from frigate.config import FrigateConfig from frigate.object_detection import LocalObjectDetector from frigate.motion import MotionDetector from frigate.object_processing import CameraState -from frigate.objects import ObjectTracker from frigate.util import ( EventsPerSecond, SharedMemoryFrameManager, @@ -107,7 +108,7 @@ class ProcessClip: motion_detector = MotionDetector(self.frame_shape, self.camera_config.motion) motion_detector.save_images = False - object_tracker = ObjectTracker(self.camera_config.detect) + object_tracker = CentroidTracker(self.camera_config.detect) process_info = { "process_fps": mp.Value("d", 0.0), "detection_fps": mp.Value("d", 0.0), @@ -247,7 +248,7 @@ def process(path, label, output, debug_path): clips.append(path) json_config = { - "mqtt": {"host": "mqtt"}, + "mqtt": {"enabled": False}, "detectors": {"coral": {"type": "edgetpu", "device": "usb"}}, "cameras": { "camera": { @@ -281,7 +282,7 @@ def process(path, label, output, debug_path): json_config["cameras"]["camera"]["ffmpeg"]["inputs"][0]["path"] = c frigate_config = FrigateConfig(**json_config) - runtime_config = frigate_config.runtime_config + runtime_config = frigate_config.runtime_config() runtime_config.cameras["camera"].create_ffmpeg_cmds() process_clip = ProcessClip(c, frame_shape, runtime_config)