main_2024.py detect target

main_2024.py

@@ -45,6 +45,7 @@ def main() -> int:
     parser = argparse.ArgumentParser()
     parser.add_argument("--cpu", action="store_true", help="option to force cpu")
     parser.add_argument("--full", action="store_true", help="option to force full precision")
+    parser.add_argument("--show-annotated", action="store_true", help="option to show annotated image")
     args = parser.parse_args()
 
     # Set constants
@@ -64,6 +65,7 @@ def main() -> int:
         DETECT_TARGET_OVERRIDE_FULL_PRECISION = args.full
         DETECT_TARGET_SAVE_NAME_PREFIX = config["detect_target"]["save_prefix"]
         DETECT_TARGET_SAVE_PREFIX = f"{LOG_DIRECTORY_PATH}/{DETECT_TARGET_SAVE_NAME_PREFIX}"
+        DETECT_TARGET_SHOW_ANNOTATED = args.show_annotated
 
         FLIGHT_INTERFACE_ADDRESS = config["flight_interface"]["address"]
         FLIGHT_INTERFACE_TIMEOUT = config["flight_interface"]["timeout"]
@@ -113,6 +115,7 @@ def main() -> int:
             DETECT_TARGET_MODEL_PATH,
             DETECT_TARGET_OVERRIDE_FULL_PRECISION,
             DETECT_TARGET_SAVE_PREFIX,
+            DETECT_TARGET_SHOW_ANNOTATED,
             video_input_to_detect_target_queue,
             detect_target_to_main_queue,
             controller,
@@ -139,9 +142,17 @@ def main() -> int:
 
     while True:
         try:
-            image = detect_target_to_main_queue.queue.get_nowait()
+            detections = detect_target_to_main_queue.queue.get_nowait()
         except queue.Empty:
-            image = None
+            detections = None
+
+        if detections is not None:
+            print("timestamp: " + str(detections.timestamp))
+            print("detections: " + str(len(detections.detections)))
+            for detection in detections.detections:
+                print("    label: " + str(detection.label))
+                print("    confidence: " + str(detection.confidence))
+            print("")
 
         odometry_and_time_info: "odometry_and_time.OdometryAndTime | None" = \
             flight_interface_to_main_queue.queue.get()
@@ -160,12 +171,8 @@ def main() -> int:
             print("pitch: " + str(orientation.pitch))
             print("")
 
-        if image is None:
-            continue
-
-        cv2.imshow("Landing Pad Detector", image)
-
-        if cv2.waitKey(1) & 0xFF == ord('q'):
+        if cv2.waitKey(1) == ord('q'):
+            print("Exiting main loop")
             break
 
     # Teardown
@@ -179,6 +186,8 @@ def main() -> int:
     detect_target_manager.join_workers()
     flight_interface_manager.join_workers()
 
+    cv2.destroyAllWindows()
+
     return 0
 
 

modules/detect_target/detect_target.py

@@ -17,7 +17,7 @@ class DetectTarget:
     """
     Contains the YOLOv8 model for prediction.
     """
-    def __init__(self, device: "str | int", model_path: str, override_full: bool, save_name: str = ""):
+    def __init__(self, device: "str | int", model_path: str, override_full: bool, show_annotations: bool = False, save_name: str = ""):
         """
         device: name of target device to run inference on (i.e. "cpu" or cuda device 0, 1, 2, 3).
         model_path: path to the YOLOv8 model.
@@ -28,17 +28,18 @@ def __init__(self, device: "str | int", model_path: str, override_full: bool, sa
         self.__model = ultralytics.YOLO(model_path)
         self.__counter = 0
         self.__enable_half_precision = False if self.__device == "cpu" else True
+        self.__show_annotations = show_annotations
         if override_full:
             self.__enable_half_precision = False
         self.__filename_prefix = ""
         if save_name != "":
             self.__filename_prefix = save_name + "_" + str(int(time.time())) + "_"
 
-    def run(self, data: image_and_time.ImageAndTime) -> "tuple[bool, np.ndarray | None]":
+    def run(self, data: image_and_time.ImageAndTime) -> "tuple[bool, detections_and_time.DetectionsAndTime | None]":
         """
         Returns annotated image.
-        TODO: Change to DetectionsAndTime
         """
+        print(self.__enable_half_precision)
         image = data.image
         predictions = self.__model.predict(
             source=image,
@@ -50,7 +51,6 @@ def run(self, data: image_and_time.ImageAndTime) -> "tuple[bool, np.ndarray | No
         if len(predictions) == 0:
             return False, None
 
-        # TODO: Change this to DetectionsAndTime for image and telemetry merge for 2024
         image_annotated = predictions[0].plot(conf=True)
 
         # Processing object detection
@@ -89,7 +89,9 @@ def run(self, data: image_and_time.ImageAndTime) -> "tuple[bool, np.ndarray | No
 
             self.__counter += 1
 
-        # TODO: Change this to DetectionsAndTime
-        return True, image_annotated
+        if self.__show_annotations:
+            cv2.imshow("Annotated", image_annotated)
+
+        return True, detections
 
 # pylint: enable=too-few-public-methods

modules/detect_target/detect_target_worker.py

@@ -13,6 +13,7 @@ def detect_target_worker(device: "str | int",
                          model_path: str,
                          override_full: bool,
                          save_name: str,
+                         show_annotations: bool,
                          input_queue: queue_proxy_wrapper.QueueProxyWrapper,
                          output_queue: queue_proxy_wrapper.QueueProxyWrapper,
                          controller: worker_controller.WorkerController):
@@ -23,7 +24,7 @@ def detect_target_worker(device: "str | int",
     input_queue and output_queue are data queues.
     controller is how the main process communicates to this worker process.
     """
-    detector = detect_target.DetectTarget(device, model_path, override_full, save_name)
+    detector = detect_target.DetectTarget(device, model_path, override_full, show_annotations, save_name)
 
     while not controller.is_exit_requested():
         controller.check_pause()

tests/model_example/bounding_box_bus.txt

@@ -0,0 +1,5 @@
+9.213867187500000000e-01 5.000000000000000000e+00 1.518750000000000000e+01 2.286562500000000000e+02 8.083125000000000000e+02 7.484062500000000000e+02
+8.930664062500000000e-01 0.000000000000000000e+00 6.674062500000000000e+02 3.898125000000000000e+02 8.091562500000000000e+02 8.783437500000000000e+02
+8.867187500000000000e-01 0.000000000000000000e+00 4.978125000000000000e+01 4.016250000000000000e+02 2.444765625000000000e+02 9.019687500000000000e+02
+8.779296875000000000e-01 0.000000000000000000e+00 2.223281250000000000e+02 4.083750000000000000e+02 3.455156250000000000e+02 8.606250000000000000e+02
+6.113281250000000000e-01 0.000000000000000000e+00 1.582031250000000000e-01 5.509687500000000000e+02 6.517968750000000000e+01 8.690625000000000000e+02

tests/model_example/bounding_box_zidane.txt

@@ -0,0 +1,3 @@
+8.896484375000000000e-01 0.000000000000000000e+00 7.480000000000000000e+02 4.150000000000000000e+01 1.140000000000000000e+03 7.130000000000000000e+02
+8.847656250000000000e-01 0.000000000000000000e+00 1.450000000000000000e+02 2.000000000000000000e+02 1.108000000000000000e+03 7.130000000000000000e+02
+7.167968750000000000e-01 2.700000000000000000e+01 4.375000000000000000e+02 4.347500000000000000e+02 5.300000000000000000e+02 7.170000000000000000e+02

tests/model_example/generate_expected.py

@@ -2,19 +2,31 @@
 Generates expected output using pretrained default model and images.
 TODO: PointsAndTime
 """
+import pathlib
 
 import cv2
+import numpy as np
 import ultralytics
 
 
+TEST_PATH = pathlib.Path("tests", "model_example")
+
 # Downloaded from: https://github.com/ultralytics/assets/releases
-MODEL_PATH = "tests/model_example/yolov8s_ultralytics_pretrained_default.pt"
+MODEL_PATH = pathlib.Path(TEST_PATH, "yolov8s_ultralytics_pretrained_default.pt")
+
+BUS_IMAGE_PATH = pathlib.Path(TEST_PATH, "bus.jpg")
+ZIDANE_IMAGE_PATH = pathlib.Path(TEST_PATH, "zidane.jpg")
+
+BUS_IMAGE_ANNOTATED_PATH = pathlib.Path(TEST_PATH, "bus_annotated.png")
+ZIDANE_IMAGE_ANNOTATED_PATH = pathlib.Path(TEST_PATH, "zidane_annotated.png")
+BUS_BOUNDING_BOX_PATH = pathlib.Path(TEST_PATH, "bounding_box_bus.txt")
+ZIDANE_BOUNDING_BOX_PATH = pathlib.Path(TEST_PATH, "bounding_box_zidane.txt")
 
 
 if __name__ == "__main__":
     model = ultralytics.YOLO(MODEL_PATH)
-    image_bus = cv2.imread("tests/model_example/bus.jpg")
-    image_zidane = cv2.imread("tests/model_example/zidane.jpg")
+    image_bus = cv2.imread(BUS_IMAGE_PATH)
+    image_zidane = cv2.imread(ZIDANE_IMAGE_PATH)
 
     # ultralytics saves as .jpg , bad for testing reproducibility
     results_bus = model.predict(
@@ -34,7 +46,25 @@
     image_zidane_annotated = results_zidane[0].plot(conf=True)
 
     # Save image
-    cv2.imwrite("tests/model_example/bus_annotated.png", image_bus_annotated)
-    cv2.imwrite("tests/model_example/zidane_annotated.png", image_zidane_annotated)
+    cv2.imwrite(BUS_IMAGE_ANNOTATED_PATH, image_bus_annotated)
+    cv2.imwrite(ZIDANE_IMAGE_ANNOTATED_PATH, image_zidane_annotated)
+
+    # Generate expected
+    bounding_box_bus = results_bus[0].boxes.xyxy.detach().cpu().numpy()
+    bounding_box_zidane = results_zidane[0].boxes.xyxy.detach().cpu().numpy()
+
+    conf_bus = results_bus[0].boxes.conf.detach().cpu().numpy()
+    conf_zidane = results_zidane[0].boxes.conf.detach().cpu().numpy()
+
+    labels_bus = results_bus[0].boxes.cls.detach().cpu().numpy()
+    labels_zidane = results_zidane[0].boxes.cls.detach().cpu().numpy()
+
+    predictions_bus = np.insert(bounding_box_bus, 0, [conf_bus, labels_bus], axis=1)
+    predictions_zidane = np.insert(bounding_box_zidane, 0, [conf_zidane, labels_zidane], axis=1)
+
+    # Save expected to text file
+    # Format: [confidence, label, x1, y1, x2, y2]
+    np.savetxt(BUS_BOUNDING_BOX_PATH, predictions_bus)
+    np.savetxt(ZIDANE_BOUNDING_BOX_PATH, predictions_zidane)
 
     print("Done!")