From 660f3b46a7f8147c923709f83798cccd839a740e Mon Sep 17 00:00:00 2001
From: zrezke <jeretina.filip@gmail.com>
Date: Tue, 27 Feb 2024 17:51:03 +0100
Subject: [PATCH] Fix thermal model, and coord norm.

---
 rerun_py/depthai_viewer/_backend/device.py         |  3 ++-
 rerun_py/depthai_viewer/_backend/packet_handler.py | 10 +++++-----
 rerun_py/depthai_viewer/install_requirements.py    |  9 ++++++++-
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/rerun_py/depthai_viewer/_backend/device.py b/rerun_py/depthai_viewer/_backend/device.py
index 40a1e1203a9a..37f4f38003de 100644
--- a/rerun_py/depthai_viewer/_backend/device.py
+++ b/rerun_py/depthai_viewer/_backend/device.py
@@ -549,10 +549,11 @@ def update_pipeline(self, runtime_only: bool) -> Message:
                 if "yolo" in config.ai_model.path:
                     yolo = self._oak.pipeline.createYoloDetectionNetwork()
                     yolo.setBlobPath(model_path)
-                    yolo.setConfidenceThreshold(0.5)
                     if "yolov6n_thermal_people_256x192" == config.ai_model.path:
+                        yolo.setConfidenceThreshold(0.5)
                         yolo.setNumClasses(1)
                         yolo.setCoordinateSize(4)
+                        yolo.setIouThreshold(0.5)
                 cam_node.raw.link(yolo.input)
                 xlink_out_yolo = self._oak.pipeline.createXLinkOut()
                 xlink_out_yolo.setStreamName("yolo")
diff --git a/rerun_py/depthai_viewer/_backend/packet_handler.py b/rerun_py/depthai_viewer/_backend/packet_handler.py
index 24e80a1dc4b5..88331d44dfa0 100644
--- a/rerun_py/depthai_viewer/_backend/packet_handler.py
+++ b/rerun_py/depthai_viewer/_backend/packet_handler.py
@@ -90,7 +90,7 @@ def _dai_detections_to_rects_colors_labels(
         colors = []
         labels = []
         for detection in packet.detections:
-            rects.append(self._rect_from_detection(detection, context.frame_width, context.frame_height))
+            rects.append(self._rect_from_detection(detection, context.frame_height, context.frame_width))
             colors.append([0, 255, 0])
             label = ""
             # Open model zoo models output label index
@@ -273,10 +273,10 @@ def _on_age_gender_packet(self, packet: TwoStagePacket, component: NNComponent)
 
     def _rect_from_detection(self, detection: dai.ImgDetection, max_height: int, max_width: int) -> List[int]:
         return [
-            int(max(min(detection.xmin, max_width), 0) * max_width),
-            int(max(min(detection.xmax, max_height), 0) * max_height),
-            int(max(min(detection.ymax, max_width), 0) * max_width),
-            int(max(min(detection.ymin, max_height), 0) * max_height),
+            int(min(max(detection.xmin, 0.0), 1.0) * max_width),
+            int(min(max(detection.ymin, 0.0), 1.0) * max_height),
+            int(min(max(detection.xmax, 0.0), 1.0) * max_width),
+            int(min(max(detection.ymax, 0.0), 1.0) * max_height),
         ]
 
 
diff --git a/rerun_py/depthai_viewer/install_requirements.py b/rerun_py/depthai_viewer/install_requirements.py
index 1a5cac475e19..c2e0bdbd66e2 100644
--- a/rerun_py/depthai_viewer/install_requirements.py
+++ b/rerun_py/depthai_viewer/install_requirements.py
@@ -57,7 +57,14 @@ def download_blobs() -> None:
     }
     for model, zoo_type in models.items():
         # With use_cache=True, blobconverter will not download / move the blob to model_dir...
-        blobconverter.from_zoo(model, zoo_type=zoo_type, shaves=6, output_dir=model_dir, use_cache=False)
+        blobconverter.from_zoo(
+            model,
+            zoo_type=zoo_type,
+            shaves=6,
+            output_dir=model_dir,
+            use_cache=False,
+            compile_params=["-ip FP16"] if "thermal" in model else None,
+        )
 
 
 def dependencies_installed() -> bool: