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,22 +1,14 @@
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:
class CentroidTracker(ObjectTracker):
def __init__(self, config: DetectConfig):
self.tracked_objects = {}
self.disappeared = {}
@ -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
]

View File

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

View File

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