diff --git a/frigate/data_processing/real_time/bird_processor.py b/frigate/data_processing/real_time/bird_processor.py index a6dfb8510..c96fb1868 100644 --- a/frigate/data_processing/real_time/bird_processor.py +++ b/frigate/data_processing/real_time/bird_processor.py @@ -3,10 +3,12 @@ import logging import os +import cv2 import numpy as np from frigate.config import FrigateConfig from frigate.const import MODEL_CACHE_DIR +from frigate.util.object import calculate_region from ..types import DataProcessorMetrics from .api import RealTimeProcessorApi @@ -75,16 +77,36 @@ class BirdProcessor(RealTimeProcessorApi): if obj_data["label"] != "bird": return - self.interpreter.set_tensor(self.tensor_input_details[0]["index"], frame) + x, y, x2, y2 = calculate_region( + frame.shape, + obj_data["box"][0], + obj_data["box"][1], + obj_data["box"][2], + obj_data["box"][3], + 224, + 1.4, + ) + + rgb = cv2.cvtColor(frame, cv2.COLOR_YUV2RGB_I420) + input = rgb[ + y:y2, + x:x2, + ] + + logger.info(f"input shape is {input.shape}") + cv2.imwrite("/media/frigate/test_class.png", input) + + input = np.expand_dims(input, axis=0) + + self.interpreter.set_tensor(self.tensor_input_details[0]["index"], input) self.interpreter.invoke() - res = self.interpreter.get_tensor(self.tensor_output_details[0]["index"])[0] - non_zero_indices = res > 0 - class_ids = np.argpartition(-res, 20)[:20] - class_ids = class_ids[np.argsort(-res[class_ids])] - class_ids = class_ids[non_zero_indices[class_ids]] - scores = res[class_ids] - boxes = np.full((scores.shape[0], 4), -1, np.float32) - count = len(scores) + res: np.ndarray = self.interpreter.get_tensor(self.tensor_output_details[0]["index"])[0] + probs = res / res.sum(axis=0) + best_id = np.argmax(probs) + score = probs[best_id] def handle_request(self, request_data): - return None \ No newline at end of file + return None + + def expire_object(self, object_id): + pass