mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-02-20 13:54:36 +01:00
Refactor processors and add LPR postprocessing (#16722)
* recordings data pub/sub * function to process recording stream frames * model runner * lpr model runner * refactor to mixin class and use model runner * separate out realtime and post processors * move model and mixin folders * basic postprocessor * clean up * docs * postprocessing logic * clean up * return none if recordings are disabled * run postprocessor handle_requests too * tweak expansion * add put endpoint * postprocessor tweaks with endpoint
This commit is contained in:
Josh Hawkins
GitHub
@@ -5,16 +5,22 @@ from abc import ABC, abstractmethod
|
||||
|
||||
from frigate.config import FrigateConfig
|
||||
|
||||
from ..types import DataProcessorMetrics, PostProcessDataEnum
|
||||
from ..types import DataProcessorMetrics, DataProcessorModelRunner, PostProcessDataEnum
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PostProcessorApi(ABC):
|
||||
@abstractmethod
|
||||
def __init__(self, config: FrigateConfig, metrics: DataProcessorMetrics) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
config: FrigateConfig,
|
||||
metrics: DataProcessorMetrics,
|
||||
model_runner: DataProcessorModelRunner,
|
||||
) -> None:
|
||||
self.config = config
|
||||
self.metrics = metrics
|
||||
self.model_runner = model_runner
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
|
||||
231
frigate/data_processing/post/license_plate.py
Normal file
231
frigate/data_processing/post/license_plate.py
Normal file
@@ -0,0 +1,231 @@
|
||||
"""Handle post processing for license plate recognition."""
|
||||
|
||||
import datetime
|
||||
import logging
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from peewee import DoesNotExist
|
||||
|
||||
from frigate.comms.embeddings_updater import EmbeddingsRequestEnum
|
||||
from frigate.config import FrigateConfig
|
||||
from frigate.data_processing.common.license_plate.mixin import (
|
||||
WRITE_DEBUG_IMAGES,
|
||||
LicensePlateProcessingMixin,
|
||||
)
|
||||
from frigate.data_processing.common.license_plate.model import (
|
||||
LicensePlateModelRunner,
|
||||
)
|
||||
from frigate.data_processing.types import PostProcessDataEnum
|
||||
from frigate.models import Recordings
|
||||
from frigate.util.image import get_image_from_recording
|
||||
|
||||
from ..types import DataProcessorMetrics
|
||||
from .api import PostProcessorApi
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LicensePlatePostProcessor(LicensePlateProcessingMixin, PostProcessorApi):
|
||||
def __init__(
|
||||
self,
|
||||
config: FrigateConfig,
|
||||
metrics: DataProcessorMetrics,
|
||||
model_runner: LicensePlateModelRunner,
|
||||
detected_license_plates: dict[str, dict[str, any]],
|
||||
):
|
||||
self.detected_license_plates = detected_license_plates
|
||||
self.model_runner = model_runner
|
||||
self.lpr_config = config.lpr
|
||||
self.config = config
|
||||
super().__init__(config, metrics, model_runner)
|
||||
|
||||
def __update_metrics(self, duration: float) -> None:
|
||||
"""
|
||||
Update inference metrics.
|
||||
"""
|
||||
self.metrics.alpr_pps.value = (self.metrics.alpr_pps.value * 9 + duration) / 10
|
||||
|
||||
def process_data(
|
||||
self, data: dict[str, any], data_type: PostProcessDataEnum
|
||||
) -> None:
|
||||
"""Look for license plates in recording stream image
|
||||
Args:
|
||||
data (dict): containing data about the input.
|
||||
data_type (enum): Describing the data that is being processed.
|
||||
|
||||
Returns:
|
||||
None.
|
||||
"""
|
||||
start = datetime.datetime.now().timestamp()
|
||||
|
||||
event_id = data["event_id"]
|
||||
camera_name = data["camera"]
|
||||
|
||||
if data_type == PostProcessDataEnum.recording:
|
||||
obj_data = data["obj_data"]
|
||||
frame_time = obj_data["frame_time"]
|
||||
recordings_available_through = data["recordings_available"]
|
||||
|
||||
if frame_time > recordings_available_through:
|
||||
logger.debug(
|
||||
f"LPR post processing: No recordings available for this frame time {frame_time}, available through {recordings_available_through}"
|
||||
)
|
||||
|
||||
elif data_type == PostProcessDataEnum.tracked_object:
|
||||
# non-functional, need to think about snapshot time
|
||||
obj_data = data["event"]["data"]
|
||||
obj_data["id"] = data["event"]["id"]
|
||||
obj_data["camera"] = data["event"]["camera"]
|
||||
# TODO: snapshot time?
|
||||
frame_time = data["event"]["start_time"]
|
||||
|
||||
else:
|
||||
logger.error("No data type passed to LPR postprocessing")
|
||||
return
|
||||
|
||||
recording_query = (
|
||||
Recordings.select(
|
||||
Recordings.path,
|
||||
Recordings.start_time,
|
||||
)
|
||||
.where(
|
||||
(
|
||||
(frame_time >= Recordings.start_time)
|
||||
& (frame_time <= Recordings.end_time)
|
||||
)
|
||||
)
|
||||
.where(Recordings.camera == camera_name)
|
||||
.order_by(Recordings.start_time.desc())
|
||||
.limit(1)
|
||||
)
|
||||
|
||||
try:
|
||||
recording: Recordings = recording_query.get()
|
||||
time_in_segment = frame_time - recording.start_time
|
||||
codec = "mjpeg"
|
||||
|
||||
image_data = get_image_from_recording(
|
||||
self.config.ffmpeg, recording.path, time_in_segment, codec, None
|
||||
)
|
||||
|
||||
if not image_data:
|
||||
logger.debug(
|
||||
"LPR post processing: Unable to fetch license plate from recording"
|
||||
)
|
||||
|
||||
# Convert bytes to numpy array
|
||||
image_array = np.frombuffer(image_data, dtype=np.uint8)
|
||||
|
||||
if len(image_array) == 0:
|
||||
logger.debug("LPR post processing: No image")
|
||||
return
|
||||
|
||||
image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
|
||||
|
||||
except DoesNotExist:
|
||||
logger.debug("Error fetching license plate for postprocessing")
|
||||
return
|
||||
|
||||
if WRITE_DEBUG_IMAGES:
|
||||
cv2.imwrite(f"debug/frames/lpr_post_{start}.jpg", image)
|
||||
|
||||
# convert to yuv for processing
|
||||
frame = cv2.cvtColor(image, cv2.COLOR_BGR2YUV_I420)
|
||||
|
||||
detect_width = self.config.cameras[camera_name].detect.width
|
||||
detect_height = self.config.cameras[camera_name].detect.height
|
||||
|
||||
# Scale the boxes based on detect dimensions
|
||||
scale_x = image.shape[1] / detect_width
|
||||
scale_y = image.shape[0] / detect_height
|
||||
|
||||
# Determine which box to enlarge based on detection mode
|
||||
if self.requires_license_plate_detection:
|
||||
# Scale and enlarge the car box
|
||||
box = obj_data.get("box")
|
||||
if not box:
|
||||
return
|
||||
|
||||
# Scale original car box to detection dimensions
|
||||
left = int(box[0] * scale_x)
|
||||
top = int(box[1] * scale_y)
|
||||
right = int(box[2] * scale_x)
|
||||
bottom = int(box[3] * scale_y)
|
||||
box = [left, top, right, bottom]
|
||||
else:
|
||||
# Get the license plate box from attributes
|
||||
if not obj_data.get("current_attributes"):
|
||||
return
|
||||
|
||||
license_plate = None
|
||||
for attr in obj_data["current_attributes"]:
|
||||
if attr.get("label") != "license_plate":
|
||||
continue
|
||||
if license_plate is None or attr.get("score", 0.0) > license_plate.get(
|
||||
"score", 0.0
|
||||
):
|
||||
license_plate = attr
|
||||
|
||||
if not license_plate or not license_plate.get("box"):
|
||||
return
|
||||
|
||||
# Scale license plate box to detection dimensions
|
||||
orig_box = license_plate["box"]
|
||||
left = int(orig_box[0] * scale_x)
|
||||
top = int(orig_box[1] * scale_y)
|
||||
right = int(orig_box[2] * scale_x)
|
||||
bottom = int(orig_box[3] * scale_y)
|
||||
box = [left, top, right, bottom]
|
||||
|
||||
width_box = right - left
|
||||
height_box = bottom - top
|
||||
|
||||
# Enlarge box slightly to account for drift in detect vs recording stream
|
||||
enlarge_factor = 0.3
|
||||
new_left = max(0, int(left - (width_box * enlarge_factor / 2)))
|
||||
new_top = max(0, int(top - (height_box * enlarge_factor / 2)))
|
||||
new_right = min(image.shape[1], int(right + (width_box * enlarge_factor / 2)))
|
||||
new_bottom = min(
|
||||
image.shape[0], int(bottom + (height_box * enlarge_factor / 2))
|
||||
)
|
||||
|
||||
keyframe_obj_data = obj_data.copy()
|
||||
if self.requires_license_plate_detection:
|
||||
# car box
|
||||
keyframe_obj_data["box"] = [new_left, new_top, new_right, new_bottom]
|
||||
else:
|
||||
# Update the license plate box in the attributes
|
||||
new_attributes = []
|
||||
for attr in obj_data["current_attributes"]:
|
||||
if attr.get("label") == "license_plate":
|
||||
new_attr = attr.copy()
|
||||
new_attr["box"] = [new_left, new_top, new_right, new_bottom]
|
||||
new_attributes.append(new_attr)
|
||||
else:
|
||||
new_attributes.append(attr)
|
||||
keyframe_obj_data["current_attributes"] = new_attributes
|
||||
|
||||
# run the frame through lpr processing
|
||||
logger.debug(f"Post processing plate: {event_id}, {frame_time}")
|
||||
self.lpr_process(keyframe_obj_data, frame)
|
||||
|
||||
self.__update_metrics(datetime.datetime.now().timestamp() - start)
|
||||
|
||||
def handle_request(self, topic, request_data) -> dict[str, any] | None:
|
||||
if topic == EmbeddingsRequestEnum.reprocess_plate.value:
|
||||
event = request_data["event"]
|
||||
|
||||
self.process_data(
|
||||
{
|
||||
"event_id": event["id"],
|
||||
"camera": event["camera"],
|
||||
"event": event,
|
||||
},
|
||||
PostProcessDataEnum.tracked_object,
|
||||
)
|
||||
|
||||
return {
|
||||
"message": "Successfully requested reprocessing of license plate.",
|
||||
"success": True,
|
||||
}
|
||||
Reference in New Issue
Block a user