Fix typing issues in vision

bit-bots · Jan 22, 2025 · 5147b60 · 5147b60
1 parent 2d9cff3
commit 5147b60
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/bitbots_vision/bitbots_vision/vision_modules/yoeo/utils.py b/bitbots_vision/bitbots_vision/vision_modules/yoeo/utils.py
@@ -5,7 +5,7 @@
 import cv2
 import numpy as np
 import rclpy
-from jaxtyping import UInt8
+from jaxtyping import Float64, UInt8
 
 logger = rclpy.logging.get_logger("yoeo_handler_utils")
 
@@ -158,10 +158,10 @@ def get_info(self) -> ImagePreProcessorData:
             padding_bottom=self._padding_bottom,
             padding_left=self._padding_left,
             padding_right=self._padding_right,
-            max_dim=np.max(self._image_dimensions_HW),
+            max_dim=int(np.max(self._image_dimensions_HW)),
         )
 
-    def process(self, image: UInt8[np.ndarray, "h w 3"]) -> UInt8[np.ndarray, "3 network_input_h network_input_w"]:
+    def process(self, image: UInt8[np.ndarray, "h w 3"]) -> Float64[np.ndarray, "3 network_input_h network_input_w"]:
         self._image_dimensions_HW = image.shape[:2]  # type: ignore[assignment]
         self._calculate_paddings()