From 58bf9f30bba3f459518f4afd25fb9ae217c02e88 Mon Sep 17 00:00:00 2001
From: zrezke <jeretina.filip@gmail.com>
Date: Mon, 18 Mar 2024 14:08:36 +0100
Subject: [PATCH] Fix bboxes and age-gender-recognition

---
 rerun_py/depthai_viewer/_backend/device.py    | 17 +++++-----
 .../depthai_viewer/_backend/packet_handler.py | 32 +++++++++++++++----
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/rerun_py/depthai_viewer/_backend/device.py b/rerun_py/depthai_viewer/_backend/device.py
index 677393690023..4456b9b5183e 100644
--- a/rerun_py/depthai_viewer/_backend/device.py
+++ b/rerun_py/depthai_viewer/_backend/device.py
@@ -6,17 +6,13 @@
 from typing import Dict, List, Optional, Tuple
 
 import depthai as dai
+import depthai_viewer as viewer
 import numpy as np
 from depthai_sdk import OakCamera
 from depthai_sdk.classes.packet_handlers import ComponentOutput
 from depthai_sdk.components import CameraComponent, NNComponent, StereoComponent
-from depthai_sdk.components.camera_helper import (
-    getClosestIspScale,
-)
+from depthai_sdk.components.camera_helper import getClosestIspScale
 from depthai_sdk.components.tof_component import Component
-from numpy.typing import NDArray
-
-import depthai_viewer as viewer
 from depthai_viewer._backend.device_configuration import (
     ALL_NEURAL_NETWORKS,
     CameraConfiguration,
@@ -40,9 +36,14 @@
     Message,
     WarningMessage,
 )
-from depthai_viewer._backend.packet_handler import DetectionContext, PacketHandler, PacketHandlerContext
+from depthai_viewer._backend.packet_handler import (
+    DetectionContext,
+    PacketHandler,
+    PacketHandlerContext,
+)
 from depthai_viewer._backend.store import Store
 from depthai_viewer.install_requirements import model_dir
+from numpy.typing import NDArray
 
 
 class XlinkStatistics:
@@ -572,7 +573,7 @@ def update_pipeline(self, runtime_only: bool) -> Message:
                     WarningMessage(f"{config.ai_model.camera} is not configured, won't create NNET.")
                 )
             elif config.ai_model.path == "age-gender-recognition-retail-0013":
-                face_detection = self._oak.create_nn(model_path, cam_component)
+                face_detection = self._oak.create_nn("face-detection-retail-0004", cam_component)
                 self._nnet = self._oak.create_nn(model_path, input=face_detection)
             else:
                 self._nnet = self._oak.create_nn(model_path, cam_component)
diff --git a/rerun_py/depthai_viewer/_backend/packet_handler.py b/rerun_py/depthai_viewer/_backend/packet_handler.py
index 88331d44dfa0..80b526b88dc2 100644
--- a/rerun_py/depthai_viewer/_backend/packet_handler.py
+++ b/rerun_py/depthai_viewer/_backend/packet_handler.py
@@ -2,26 +2,32 @@
 
 import cv2
 import depthai as dai
+import depthai_viewer as viewer
 import numpy as np
 from ahrs.filters import Mahony
 from depthai_sdk.classes.packets import (  # PointcloudPacket,
     BasePacket,
+    BoundingBox,
     DepthPacket,
+    Detection,
     DetectionPacket,
     DisparityDepthPacket,
     FramePacket,
     IMUPacket,
     TwoStagePacket,
 )
-from depthai_sdk.components import CameraComponent, Component, NNComponent, StereoComponent
+from depthai_sdk.components import (
+    CameraComponent,
+    Component,
+    NNComponent,
+    StereoComponent,
+)
 from depthai_sdk.components.tof_component import ToFComponent
-from numpy.typing import NDArray
-from pydantic import BaseModel
-
-import depthai_viewer as viewer
 from depthai_viewer._backend.store import Store
 from depthai_viewer._backend.topic import Topic
 from depthai_viewer.components.rect2d import RectFormat
+from numpy.typing import NDArray
+from pydantic import BaseModel
 
 
 class PacketHandlerContext(BaseModel):  # type: ignore[misc]
@@ -234,6 +240,18 @@ def _on_detections(self, packet: DetectionPacket, component: NNComponent) -> Non
             labels=labels,
         )
 
+    def _rect_from_sdk_detection(
+        self, packet_bbox: BoundingBox, detection: Detection, max_height: int, max_width: int
+    ) -> List[int]:
+        bbox = packet_bbox.get_relative_bbox(detection.bbox)
+        (x1, y1), (x2, y2) = bbox.denormalize((max_height, max_width))
+        return [
+            max(x1, 0),
+            max(y1, 0),
+            min(x2, max_width),
+            min(y2, max_height),
+        ]
+
     def _detections_to_rects_colors_labels(
         self, packet: DetectionPacket, omz_labels: Optional[List[str]] = None
     ) -> Tuple[List[List[int]], List[List[int]], List[str]]:
@@ -242,7 +260,7 @@ def _detections_to_rects_colors_labels(
         labels = []
         for detection in packet.detections:
             rects.append(
-                self._rect_from_detection(detection.img_detection, packet.frame.shape[0], packet.frame.shape[1])
+                self._rect_from_sdk_detection(packet.bbox, detection, packet.frame.shape[0], packet.frame.shape[1])
             )
             colors.append([0, 255, 0])
             label: str = detection.label_str
@@ -265,7 +283,7 @@ def _on_age_gender_packet(self, packet: TwoStagePacket, component: NNComponent)
             cam = "color_cam" if component._get_camera_comp().is_color() else "mono_cam"
             viewer.log_rect(
                 f"{component._get_camera_comp()._socket.name}/transform/{cam}/Detection",
-                self._rect_from_detection(det.img_detection, packet.frame.shape[0], packet.frame.shape[1]),
+                self._rect_from_sdk_detection(packet.bbox, det, packet.frame.shape[0], packet.frame.shape[1]),
                 rect_format=RectFormat.XYXY,
                 color=color,
                 label=label,