fixed image annotation and modified test cases

main_2024.py

@@ -133,7 +133,7 @@ def main() -> int:
 
     while True:
         try:
-            detections = detect_target_to_main_queue.queue.get_nowait()
+            detections, image = detect_target_to_main_queue.queue.get_nowait()
         except queue.Empty:
             detections = None
 
@@ -152,7 +152,7 @@ def main() -> int:
         if detections is None:
             continue
 
-        #cv2.imshow("Landing Pad Detector", image)
+        cv2.imshow("Landing Pad Detector", image)
 
         if cv2.waitKey(1) & 0xFF == ord('q'):
             break

modules/detect_target/detect_target.py

@@ -35,7 +35,7 @@ def __init__(self, device: "str | int", model_path: str, override_full: bool, sa
         if save_name != "":
             self.__filename_prefix = save_name + "_" + str(int(time.time())) + "_"
 
-    def run(self, data: image_and_time.ImageAndTime) -> "tuple[bool, detections_and_time.DetectionsAndTime | None]":
+    def run(self, data: image_and_time.ImageAndTime) -> "tuple[bool, tuple[detections_and_time.DetectionsAndTime | None, np.ndarray | None]]":
         """
         Returns annotated image.
         """
@@ -84,8 +84,8 @@ def run(self, data: image_and_time.ImageAndTime) -> "tuple[bool, detections_and_
             self.__counter += 1
 
         if self.__show_annotations:
-            cv2.imshow("Annotated Image", image_annotated)
+            return True, (detections, image_annotated)
 
-        return True, detections
+        return True, (detections, None)
 
 # pylint: enable=too-few-public-methods

tests/test_detect_target.py

@@ -7,11 +7,9 @@
 import numpy as np
 import pytest
 import torch
-import ultralytics
 
 from modules.detect_target import detect_target
 from modules import image_and_time
-from modules import detections_and_time
 
 
 DEVICE =                        0 if torch.cuda.is_available() else "cpu"
@@ -22,14 +20,13 @@
 IMAGE_ZIDANE_PATH =             "tests/model_example/zidane.jpg"
 IMAGE_ZIDANE_ANNOTATED_PATH =   "tests/model_example/zidane_annotated.png"
 
-model = ultralytics.YOLO(MODEL_PATH)
 
 @pytest.fixture()
 def detector():
     """
     Construct DetectTarget.
     """
-    detection = detect_target.DetectTarget(DEVICE, MODEL_PATH, OVERRIDE_FULL)
+    detection = detect_target.DetectTarget(DEVICE, MODEL_PATH, OVERRIDE_FULL, show_annotations=True)
     yield detection
 
 
@@ -97,31 +94,19 @@ def test_single_bus_image(self,
         Bus image.
         """
         # Setup
-        image = cv2.imread(IMAGE_BUS_PATH)
-        prediction = model.predict(
-            source=image,
-            half=True,
-            stream=False,
-        )
-
-        boxes = prediction[0].boxes
-        expected = boxes.xyxy.detach().cpu().numpy()
+        expected = cv2.imread(IMAGE_BUS_ANNOTATED_PATH)
         assert expected is not None
 
         # Run
-        result, actual = detector.run(image_bus)
-        detections = actual.detections
+        result, (detections, actual) = detector.run(image_bus)
 
         # Test
         assert result
         assert actual is not None
         assert detections is not None
 
-        error = 0
-
-        for i in range(0, len(detections)):
-            error += rmse([detections[i].x1, detections[i].y1, detections[i].x2, detections[i].y2], expected[i])
-        assert (error / len(detections)) < self.__IMAGE_DIFFERENCE_TOLERANCE
+        error = rmse(actual, expected)
+        assert error < self.__IMAGE_DIFFERENCE_TOLERANCE
 
     def test_single_zidane_image(self,
                                  detector: detect_target.DetectTarget,
@@ -130,31 +115,19 @@ def test_single_zidane_image(self,
         Zidane image.
         """
         # Setup
-        image = cv2.imread(IMAGE_ZIDANE_PATH)
-        prediction = model.predict(
-            source=image,
-            half=True,
-            stream=False,
-        )
-
-        boxes = prediction[0].boxes
-        expected = boxes.xyxy.detach().cpu().numpy()
+        expected = cv2.imread(IMAGE_ZIDANE_ANNOTATED_PATH)
         assert expected is not None
 
         # Run
-        result, actual = detector.run(image_zidane)
-        detections = actual.detections
+        result, (detections, actual) = detector.run(image_zidane)
 
         # Test
         assert result
         assert actual is not None
         assert detections is not None
 
-        error = 0
-
-        for i in range(0, len(detections)):
-            error += rmse([detections[i].x1, detections[i].y1, detections[i].x2, detections[i].y2], expected[i])
-        assert (error / len(detections)) < self.__IMAGE_DIFFERENCE_TOLERANCE
+        error = rmse(actual, expected)
+        assert error < self.__IMAGE_DIFFERENCE_TOLERANCE
 
     def test_multiple_zidane_image(self,
                                    detector: detect_target.DetectTarget,
@@ -165,15 +138,7 @@ def test_multiple_zidane_image(self,
         IMAGE_COUNT = 4
 
         # Setup
-        image = cv2.imread(IMAGE_ZIDANE_PATH)
-        prediction = model.predict(
-            source=image,
-            half=True,
-            stream=False,
-        )
-
-        boxes = prediction[0].boxes
-        expected = boxes.xyxy.detach().cpu().numpy()
+        expected = cv2.imread(IMAGE_ZIDANE_ANNOTATED_PATH)
         assert expected is not None
 
         input_images = []
@@ -189,17 +154,11 @@ def test_multiple_zidane_image(self,
 
         # Test
         for i in range(0, IMAGE_COUNT):
-            output: "tuple[bool, detections_and_time.DetectionsAndTime | None]" = outputs[i]
-            result, actual = output
-
-            detections = actual.detections
-
+            output: "tuple[bool, np.ndarray | None]" = outputs[i]
+            result, (detections, actual) = output
             assert result
             assert actual is not None
             assert detections is not None
 
-            error = 0
-
-            for i in range(0, len(detections)):
-                error += rmse([detections[i].x1, detections[i].y1, detections[i].x2, detections[i].y2], expected[i])
-            assert (error / len(detections)) < self.__IMAGE_DIFFERENCE_TOLERANCE
+            error = rmse(actual, expected)
+            assert error < self.__IMAGE_DIFFERENCE_TOLERANCE