Mirrors/blakeblackshear.frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-03-07 02:18:07 +01:00
Code Issues Projects Releases Wiki Activity

Refactor TensorRT (#18643)

Browse Source 
* Combine base and arm trt detectors

* Remove unused deps for amd64 build

* Add missing packages and cleanup ldconfig

* Expand packages for tensorflow model training

* Cleanup

* Refactor training to not reserve memory
This commit is contained in:
Nicolas Mowen
2025-06-09 08:25:33 -06:00
committed by Blake Blackshear
parent 40ab7d6c38
commit 4b57e5e265
6 changed files with 68 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
frigate/data_processing/real_time/custom_classification.py
View File

@@ -3,7 +3,6 @@
import datetime
import logging
import os
import threading
from typing import Any
import cv2
@@ -17,10 +16,8 @@ from frigate.comms.event_metadata_updater import (
from frigate.comms.inter_process import InterProcessRequestor
from frigate.config import FrigateConfig
from frigate.config.classification import CustomClassificationConfig
from frigate.const import CLIPS_DIR, MODEL_CACHE_DIR, UPDATE_MODEL_STATE
from frigate.types import ModelStatusTypesEnum
from frigate.const import CLIPS_DIR, MODEL_CACHE_DIR
from frigate.util.builtin import EventsPerSecond, InferenceSpeed, load_labels
from frigate.util.classification import train_classification_model
from frigate.util.object import box_overlaps, calculate_region
from ..types import DataProcessorMetrics
@@ -72,18 +69,6 @@ class CustomStateClassificationProcessor(RealTimeProcessorApi):
)
self.classifications_per_second.start()
def __retrain_model(self) -> None:
train_classification_model(self.model_config.name)
self.__build_detector()
self.requestor.send_data(
UPDATE_MODEL_STATE,
{
"model": self.model_config.name,
"state": ModelStatusTypesEnum.complete,
},
)
logger.info(f"Successfully loaded updated model for {self.model_config.name}")
def __update_metrics(self, duration: float) -> None:
self.classifications_per_second.update()
self.inference_speed.update(duration)
@@ -172,19 +157,15 @@ class CustomStateClassificationProcessor(RealTimeProcessorApi):
)
def handle_request(self, topic, request_data):
if topic == EmbeddingsRequestEnum.train_classification.value:
if topic == EmbeddingsRequestEnum.reload_classification_model.value:
if request_data.get("model_name") == self.model_config.name:
self.requestor.send_data(
UPDATE_MODEL_STATE,
{
"model": self.model_config.name,
"state": ModelStatusTypesEnum.training,
},
self.__build_detector()
logger.info(
f"Successfully loaded updated model for {self.model_config.name}"
)
threading.Thread(target=self.__retrain_model).start()
return {
"success": True,
"message": f"Began training {self.model_config.name} model.",
"message": f"Loaded {self.model_config.name} model.",
}
else:
return None
@@ -232,18 +213,6 @@ class CustomObjectClassificationProcessor(RealTimeProcessorApi):
prefill=0,
)
def __retrain_model(self) -> None:
train_classification_model(self.model_config.name)
self.__build_detector()
self.requestor.send_data(
UPDATE_MODEL_STATE,
{
"model": self.model_config.name,
"state": ModelStatusTypesEnum.complete,
},
)
logger.info(f"Successfully loaded updated model for {self.model_config.name}")
def __update_metrics(self, duration: float) -> None:
self.classifications_per_second.update()
self.inference_speed.update(duration)
@@ -307,19 +276,14 @@ class CustomObjectClassificationProcessor(RealTimeProcessorApi):
self.detected_objects[obj_data["id"]] = score
def handle_request(self, topic, request_data):
if topic == EmbeddingsRequestEnum.train_classification.value:
if topic == EmbeddingsRequestEnum.reload_classification_model.value:
if request_data.get("model_name") == self.model_config.name:
self.requestor.send_data(
UPDATE_MODEL_STATE,
{
"model": self.model_config.name,
"state": ModelStatusTypesEnum.training,
},
logger.info(
f"Successfully loaded updated model for {self.model_config.name}"
)
threading.Thread(target=self.__retrain_model).start()
return {
"success": True,
"message": f"Began training {self.model_config.name} model.",
"message": f"Loaded {self.model_config.name} model.",
}
else:
return None

6
frigate/data_processing/types.py
View File

@@ -20,8 +20,8 @@ class DataProcessorMetrics:
alpr_pps: Synchronized
yolov9_lpr_speed: Synchronized
yolov9_lpr_pps: Synchronized
classification_speeds: dict[str, Synchronized]
classification_cps: dict[str, Synchronized]
classification_speeds: dict[str, Synchronized] = {}
classification_cps: dict[str, Synchronized] = {}
def __init__(self, custom_classification_models: list[str]):
self.image_embeddings_speed = mp.Value("d", 0.0)
@@ -36,8 +36,6 @@ class DataProcessorMetrics:
self.yolov9_lpr_pps = mp.Value("d", 0.0)
if custom_classification_models:
self.classification_speeds = {}
self.classification_cps = {}
for key in custom_classification_models:
self.classification_speeds[key] = mp.Value("d", 0.0)
self.classification_cps[key] = mp.Value("d", 0.0)