object detection

2025-05-16 15:22:23 -04:00
parent 81a4182359
commit fe7f27a47a
1 changed files with 63 additions and 82 deletions
@@ -1,4 +1,4 @@
-from typing import List
+from typing import Any, List

 import mediapipe as mp
 import numpy as np
@@ -8,12 +8,11 @@ from mediapipe.tasks.python import vision
 from mediapipe.tasks.python.core.base_options import BaseOptions

 from ...utils.timing import TimestampProvider
-
-# Import new types
+from ..common import BaseDetector
 from ..types import BoundingBox, ObjectDetectionCategory, ObjectDetectionResult


-class ObjectDetector:
+class ObjectDetector(BaseDetector[ObjectDetectionResult]):
    """Detects objects using MediaPipe ObjectDetector."""

    def __init__(self, model_path: str):
@@ -30,22 +29,61 @@ class ObjectDetector:
        self._current_options = None
        self._timestamp_provider = None  # Added

-    def _initialize_detector(self, score_threshold: float, max_results: int, running_mode: str, delegate: str):  # Added running_mode
-        """Initializes the ObjectDetector."""
-        mode_enum = vision.RunningMode.IMAGE if running_mode == "image" else vision.RunningMode.VIDEO  # Added
-        delegate_enum = BaseOptions.Delegate.CPU if delegate.lower() == "cpu" else BaseOptions.Delegate.GPU
-        base_options = python.BaseOptions(model_asset_path=self.model_path, delegate=delegate_enum)
-        options = vision.ObjectDetectorOptions(
+    def _create_detector_options(self, base_options: python.BaseOptions,
+                               mode_enum: vision.RunningMode, **kwargs) -> vision.ObjectDetectorOptions:
+        """Create ObjectDetector-specific options with parameters:
+            - score_threshold: Minimum score for object detection
+            - max_results: Maximum number of detection results
+        """
+        return vision.ObjectDetectorOptions(
            base_options=base_options,
-            running_mode=mode_enum,  # Use enum mode
-            score_threshold=score_threshold,
-            max_results=max_results,
+            running_mode=mode_enum,
+            score_threshold=kwargs.get('score_threshold', 0.5),
+            max_results=kwargs.get('max_results', 5),
        )
-        self._current_options = options
-        # Only create timestamp provider for VIDEO mode
-        self._timestamp_provider = TimestampProvider() if mode_enum == vision.RunningMode.VIDEO else None
+
+    def _create_detector_instance(self, options: vision.ObjectDetectorOptions) -> vision.ObjectDetector:
        return vision.ObjectDetector.create_from_options(options)

+    def _get_options_tuple(self, running_mode: str = None, delegate: str = None, **kwargs) -> tuple:
+        if self._current_options:
+            return (
+                self._current_options.score_threshold,
+                self._current_options.max_results,
+                running_mode,
+                delegate,
+            )
+        else:
+            return (
+                kwargs.get('score_threshold', 0.5),
+                kwargs.get('max_results', 5),
+                running_mode,
+                delegate,
+            )
+
+    def _process_detection_result(self, detection_result: Any) -> List[ObjectDetectionResult]:
+        current_image_detections = []
+        if detection_result and detection_result.detections:
+            for detection in detection_result.detections:
+                bbox_mp = detection.bounding_box
+                categories_mp = detection.categories
+
+                bbox = BoundingBox(origin_x=bbox_mp.origin_x, origin_y=bbox_mp.origin_y, width=bbox_mp.width, height=bbox_mp.height)
+
+                categories = [
+                    ObjectDetectionCategory(
+                        index=cat.index,
+                        score=cat.score,
+                        display_name=cat.display_name,
+                        category_name=cat.category_name,
+                    )
+                    for cat in categories_mp
+                ]
+
+                current_image_detections.append(ObjectDetectionResult(bounding_box=bbox, categories=categories))
+
+        return current_image_detections
+
    def detect(
        self,
        image: torch.Tensor,
@@ -53,69 +91,12 @@ class ObjectDetector:
        max_results: int = 5,
        running_mode: str = "video",
        delegate: str = "cpu",
-    ) -> List[List[ObjectDetectionResult]]:  # Added running_mode
-        """Detects objects in the input image tensor.
-
-        Returns:
-            A list (batch) of lists (detections per image) of ObjectDetectionResult objects.
-        """
-        if image.dim() != 4:
-            raise ValueError("Input tensor must be in BHWC format.")
-
-        # Add running_mode to tuple check
-        new_options_tuple = (score_threshold, max_results, running_mode, delegate.lower())
-        current_options_tuple = None
-        if self._current_options:
-            current_options_tuple = (
-                self._current_options.score_threshold,
-                self._current_options.max_results,
-                "image" if self._current_options.running_mode == vision.RunningMode.IMAGE else "video",  # Map enum back to string
-                "cpu" if self._current_options.base_options.delegate == BaseOptions.Delegate.CPU else "gpu",
-            )
-
-        if self._detector_instance is None or new_options_tuple != current_options_tuple:
-            if self._detector_instance and hasattr(self._detector_instance, "close"):
-                self._detector_instance.close()
-            self._detector_instance = self._initialize_detector(score_threshold, max_results, running_mode, delegate)
-
-        batch_size = image.shape[0]
-        batch_results = []
-        # Determine the correct detection function based on mode
-        is_video_mode = self._current_options.running_mode == vision.RunningMode.VIDEO
-        detect_func = self._detector_instance.detect_for_video if is_video_mode else self._detector_instance.detect
-
-        for i in range(batch_size):
-            img_tensor = image[i]  # HWC
-            np_image = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
-            mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_image)
-
-            if is_video_mode:
-                timestamp_ms = self._timestamp_provider.next()
-                detection_result = detect_func(mp_image, timestamp_ms)
-            else:
-                detection_result = detect_func(mp_image)
-
-            current_image_detections = []
-            if detection_result and detection_result.detections:  # Added check for detection_result
-                for detection in detection_result.detections:
-                    bbox_mp = detection.bounding_box
-                    categories_mp = detection.categories
-
-                    bbox = BoundingBox(origin_x=bbox_mp.origin_x, origin_y=bbox_mp.origin_y, width=bbox_mp.width, height=bbox_mp.height)
-
-                    categories = [
-                        ObjectDetectionCategory(
-                            index=cat.index,
-                            score=cat.score,
-                            display_name=cat.display_name,
-                            category_name=cat.category_name,
-                        )
-                        for cat in categories_mp
-                    ]
-
-                    current_image_detections.append(ObjectDetectionResult(bounding_box=bbox, categories=categories))
-
-            batch_results.append(current_image_detections)
-
-        # self._detector_instance.close() # Close handled on re-init
-        return batch_results
+    ) -> List[List[ObjectDetectionResult]]:
+        """Detects objects in the input image tensor."""
+        return super().detect(
+            image,
+            running_mode=running_mode,
+            delegate=delegate,
+            score_threshold=score_threshold,
+            max_results=max_results,
+        )