refactor centroidtracker into standard API

This commit is contained in:
Blake Blackshear 2023-05-14 08:11:12 -05:00
parent 846a180a7b
commit 21778552c1
4 changed files with 35 additions and 28 deletions

13
frigate/track/__init__.py Normal file
View File

@ -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

View File

@ -1,23 +1,15 @@
import copy from collections import defaultdict
import datetime
import itertools
import multiprocessing as mp
import random import random
import string import string
import threading
import time
from collections import defaultdict
import cv2
import numpy as np import numpy as np
from scipy.spatial import distance as dist from scipy.spatial import distance as dist
from frigate.config import DetectConfig from frigate.config import DetectConfig
from frigate.track import ObjectTracker
from frigate.util import intersection_over_union from frigate.util import intersection_over_union
class CentroidTracker(ObjectTracker):
class ObjectTracker: def __init__(self, config: DetectConfig):
def __init__(self, config: DetectConfig):
self.tracked_objects = {} self.tracked_objects = {}
self.disappeared = {} self.disappeared = {}
self.positions = {} self.positions = {}
@ -141,11 +133,11 @@ class ObjectTracker:
if self.is_expired(id): if self.is_expired(id):
self.deregister(id) self.deregister(id)
def match_and_update(self, frame_time, new_objects): def match_and_update(self, frame_time, detections):
# group by name # group by name
new_object_groups = defaultdict(lambda: []) detection_groups = defaultdict(lambda: [])
for obj in new_objects: for obj in detections:
new_object_groups[obj[0]].append( detection_groups[obj[0]].append(
{ {
"label": obj[0], "label": obj[0],
"score": obj[1], "score": obj[1],
@ -160,17 +152,17 @@ class ObjectTracker:
# update any tracked objects with labels that are not # update any tracked objects with labels that are not
# seen in the current objects and deregister if needed # seen in the current objects and deregister if needed
for obj in list(self.tracked_objects.values()): 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: if self.disappeared[obj["id"]] >= self.max_disappeared:
self.deregister(obj["id"]) self.deregister(obj["id"])
else: else:
self.disappeared[obj["id"]] += 1 self.disappeared[obj["id"]] += 1
if len(new_objects) == 0: if len(detections) == 0:
return return
# track objects for each label type # track objects for each label type
for label, group in new_object_groups.items(): for label, group in detection_groups.items():
current_objects = [ current_objects = [
o for o in self.tracked_objects.values() if o["label"] == label o for o in self.tracked_objects.values() if o["label"] == label
] ]

View File

@ -19,7 +19,8 @@ from frigate.const import CACHE_DIR
from frigate.object_detection import RemoteObjectDetector from frigate.object_detection import RemoteObjectDetector
from frigate.log import LogPipe from frigate.log import LogPipe
from frigate.motion import MotionDetector 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 ( from frigate.util import (
EventsPerSecond, EventsPerSecond,
FrameManager, FrameManager,
@ -472,7 +473,7 @@ def track_camera(
name, labelmap, detection_queue, result_connection, model_config, stop_event name, labelmap, detection_queue, result_connection, model_config, stop_event
) )
object_tracker = ObjectTracker(config.detect) object_tracker = CentroidTracker(config.detect)
frame_manager = SharedMemoryFrameManager() frame_manager = SharedMemoryFrameManager()

View File

@ -1,7 +1,9 @@
import sys import sys
from typing_extensions import runtime 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 json
import logging import logging
@ -19,7 +21,6 @@ from frigate.config import FrigateConfig
from frigate.object_detection import LocalObjectDetector from frigate.object_detection import LocalObjectDetector
from frigate.motion import MotionDetector from frigate.motion import MotionDetector
from frigate.object_processing import CameraState from frigate.object_processing import CameraState
from frigate.objects import ObjectTracker
from frigate.util import ( from frigate.util import (
EventsPerSecond, EventsPerSecond,
SharedMemoryFrameManager, SharedMemoryFrameManager,
@ -107,7 +108,7 @@ class ProcessClip:
motion_detector = MotionDetector(self.frame_shape, self.camera_config.motion) motion_detector = MotionDetector(self.frame_shape, self.camera_config.motion)
motion_detector.save_images = False motion_detector.save_images = False
object_tracker = ObjectTracker(self.camera_config.detect) object_tracker = CentroidTracker(self.camera_config.detect)
process_info = { process_info = {
"process_fps": mp.Value("d", 0.0), "process_fps": mp.Value("d", 0.0),
"detection_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) clips.append(path)
json_config = { json_config = {
"mqtt": {"host": "mqtt"}, "mqtt": {"enabled": False},
"detectors": {"coral": {"type": "edgetpu", "device": "usb"}}, "detectors": {"coral": {"type": "edgetpu", "device": "usb"}},
"cameras": { "cameras": {
"camera": { "camera": {
@ -281,7 +282,7 @@ def process(path, label, output, debug_path):
json_config["cameras"]["camera"]["ffmpeg"]["inputs"][0]["path"] = c json_config["cameras"]["camera"]["ffmpeg"]["inputs"][0]["path"] = c
frigate_config = FrigateConfig(**json_config) frigate_config = FrigateConfig(**json_config)
runtime_config = frigate_config.runtime_config runtime_config = frigate_config.runtime_config()
runtime_config.cameras["camera"].create_ffmpeg_cmds() runtime_config.cameras["camera"].create_ffmpeg_cmds()
process_clip = ProcessClip(c, frame_shape, runtime_config) process_clip = ProcessClip(c, frame_shape, runtime_config)