main_2024.py detect target (#156)

* added annotated image option

* integrated in main_2024

* added cli option

* added window closing

* fixed unit tests

* fixed image annotation and modified test cases

* converted unit tests to use bounding boxes

* added device to predict arg

* added annotated image option

* integrated in main_2024

* added cli option

* added window closing

* fixed unit tests

* fixed image annotation and modified test cases

* converted unit tests to use bounding boxes

* added device to predict arg

* remove annotated image from function return

* added hardcoded values for testing target detection

* Changed tests to use absolute error tolerance, Added checks for number of detections, Changed variable names

* edited paths to use pathlib

* Fixed tests to use detections and time object

* fixed path errors in build

* Offloaded fixture setup to a separate function and added comments

* Cleaned up test_detect_target and generate_expected

* Fixed formatting

* Fixed renamed variables

* Changed tolerance

* fixed create detections

* Set tolerance to 0 and fixed formatting

* Changed formatting mostly multiline

* Final formatting changes, Updated test_detect_target_worker to use detections and time

* Some refactoring of main_2024 and final formatting changes

main_2024.py

@@ -9,6 +9,8 @@
 import cv2
 import yaml
 
+# Used in type annotation of flight interface output
+# pylint: disable-next=unused-import
 from modules import odometry_and_time
 from modules.detect_target import detect_target_worker
 from modules.flight_interface import flight_interface_worker
@@ -18,10 +20,10 @@
 from utilities.workers import worker_manager
 
 
-
 CONFIG_FILE_PATH = pathlib.Path("config.yaml")
 
-
+# Main Function
+# pylint: disable-next=too-many-locals,too-many-statements
 def main() -> int:
     """
     Main function for airside code.
@@ -45,6 +47,11 @@ 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 +71,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"]
@@ -112,6 +120,7 @@ def main() -> int:
             DETECT_TARGET_DEVICE,
             DETECT_TARGET_MODEL_PATH,
             DETECT_TARGET_OVERRIDE_FULL_PRECISION,
+            DETECT_TARGET_SHOW_ANNOTATED,
             DETECT_TARGET_SAVE_PREFIX,
             video_input_to_detect_target_queue,
             detect_target_to_main_queue,
@@ -139,9 +148,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 +177,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 +192,8 @@ def main() -> int:
     detect_target_manager.join_workers()
     flight_interface_manager.join_workers()
 
+    cv2.destroyAllWindows()
+
     return 0
 
 

modules/detect_target/detect_target.py

@@ -4,7 +4,6 @@
 import time
 
 import cv2
-import numpy as np  # TODO: Remove
 import ultralytics
 
 from .. import image_and_time
@@ -17,27 +16,43 @@ class DetectTarget:
     """
     Contains the YOLOv8 model for prediction.
     """
-    def __init__(self, device: "str | int", model_path: str, override_full: bool, save_name: str = ""):
+    # Required for logging
+    # pylint: disable-next=too-many-arguments
+    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.
         override_full: Force full precision floating point calculations.
+        show_annotations: Display annotated images.
         save_name: filename prefix for logging detections and annotated images.
         """
         self.__device = device
         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]":
+    # Required for logging
+    # pylint: disable-next=too-many-locals
+    def run(self,
+            data: image_and_time.ImageAndTime) \
+        -> "tuple[bool, detections_and_time.DetectionsAndTime | None]":
         """
-        Returns annotated image.
-        TODO: Change to DetectionsAndTime
+        Runs object detection on the provided image and returns the detections.
+
+        data: Image with a timestamp.
+
+        Return: Success and the detections.
         """
         image = data.image
         predictions = self.__model.predict(
@@ -50,7 +65,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
@@ -64,6 +78,7 @@ def run(self, data: image_and_time.ImageAndTime) -> "tuple[bool, np.ndarray | No
         if not result:
             return False, None
 
+        # Get Pylance to stop complaining
         assert detections is not None
 
         for i in range(0, boxes.shape[0]):
@@ -80,7 +95,7 @@ def run(self, data: image_and_time.ImageAndTime) -> "tuple[bool, np.ndarray | No
             filename = self.__filename_prefix + str(self.__counter)
 
             # Object detections
-            with open(filename + ".txt", "w") as file:
+            with open(filename + ".txt", "w", encoding="utf-8") as file:
                 # Use internal string representation
                 file.write(repr(detections))
 
@@ -89,7 +104,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

@@ -12,18 +12,25 @@
 def detect_target_worker(device: "str | int",
                          model_path: str,
                          override_full: bool,
+                         show_annotations: bool,
                          save_name: str,
                          input_queue: queue_proxy_wrapper.QueueProxyWrapper,
                          output_queue: queue_proxy_wrapper.QueueProxyWrapper,
                          controller: worker_controller.WorkerController):
     """
     Worker process.
 
-    device, model_path, override_full, and save_name are initial settings.
+    device, model_path, override_full, show_annotations, and save_name are initial settings.
     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

@@ -1,22 +1,33 @@
 """
 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
+    # Ultralytics saves as .jpg , bad for testing reproducibility
     results_bus = model.predict(
             source=image_bus,
             half=True,
@@ -34,7 +45,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!")