frigate/frigate/data_processing/post/review_descriptions.py

"""Post processor for review items to get descriptions."""

import logging
import shutil

import cv2

from frigate.config import FrigateConfig
from frigate.data_processing.types import PostProcessDataEnum
from frigate.genai import GenAIClient

from ..post.api import PostProcessorApi

logger = logging.getLogger(__name__)


class ReviewDescriptionProcessor(PostProcessorApi):
    def __init__(self, config: FrigateConfig, metrics, client: GenAIClient):
        super().__init__(config, metrics, None)
        self.tracked_review_items: dict[str, list[tuple[int, bytes]]] = {}
        self.genai_client = client

    def process_data(self, data, data_type):
        if data_type != PostProcessDataEnum.review:
            return

        logger.info(f"processing review {data['type']} on {data['after']['camera']}")

        id = data["after"]["id"]

        if data["type"] == "new" or data["type"] == "update":
            if id not in self.tracked_review_items:
                self.tracked_review_items[id] = []

            thumb_time = data["after"]["data"]["thumb_time"]
            thumb_path = data["after"]["thumb_path"]

            if thumb_time and thumb_path:
                if (
                    len(self.tracked_review_items[id]) > 0
                    and self.tracked_review_items[id][0] == thumb_time
                ):
                    # we have already processed this thumbnail
                    return

                shutil.copy(thumb_path, f"/media/frigate/frames/{thumb_time}.webp")
                thumb_data = cv2.imread(thumb_path)
                ret, jpg = cv2.imencode(
                    ".jpg", thumb_data, [int(cv2.IMWRITE_JPEG_QUALITY), 100]
                )

                if ret:
                    self.tracked_review_items[id].append((thumb_time, jpg.tobytes()))
        else:
            if id not in self.tracked_review_items:
                return

            final_data = data["after"]
            camera = final_data["camera"]

            if (
                data["type"] == "alert"
                and not self.config.cameras[camera].review.genai.alerts
            ):
                self.tracked_review_items.pop(id)
                return
            elif (
                data["type"] == "detection"
                and not self.config.cameras[camera].review.detections
            ):
                self.tracked_review_items.pop(id)
                return

            self.genai_client.generate_review_description(
                {
                    "camera": camera,
                    "objects": final_data["data"]["objects"],
                    "recognized_objects": final_data["data"]["sub_labels"],
                    "zones": final_data["data"]["zones"],
                },
                [r[1] for r in self.tracked_review_items[id]],
            )
            self.tracked_review_items.pop(id)

    def handle_request(self, request_data):
        pass