main_2024.py detect target

main_2024.py

@@ -15,6 +15,7 @@
 from modules.detect_target import detect_target_worker
 from modules.flight_interface import flight_interface_worker
 from modules.video_input import video_input_worker
+from modules.detections_and_time import DetectionsAndTime
 
 
 CONFIG_FILE_PATH = pathlib.Path("config.yaml")
@@ -43,6 +44,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
@@ -58,6 +60,7 @@ def main() -> int:
         DETECT_TARGET_MODEL_PATH = config["detect_target"]["model_path"]
         DETECT_TARGET_OVERRIDE_FULL_PRECISION = args.full
         DETECT_TARGET_SAVE_PREFIX = config["detect_target"]["save_prefix"]
+        DETECT_TARGET_ANNOTATE = args.show_annotated
 
         FLIGHT_INTERFACE_ADDRESS = config["flight_interface"]["address"]
         FLIGHT_INTERFACE_WORKER_PERIOD = config["flight_interface"]["worker_period"]
@@ -104,6 +107,7 @@ def main() -> int:
             DETECT_TARGET_MODEL_PATH,
             DETECT_TARGET_OVERRIDE_FULL_PRECISION,
             DETECT_TARGET_SAVE_PREFIX,
+            DETECT_TARGET_ANNOTATE,
             video_input_to_detect_target_queue,
             detect_target_to_main_queue,
             controller,
@@ -129,9 +133,9 @@ 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
 
         odometry_and_time = flight_interface_to_main_queue.queue.get()
 
@@ -145,10 +149,10 @@ def main() -> int:
             print("pitch: " + str(odometry_and_time.odometry_data.orientation.pitch))
             print("")
 
-        if image is None:
+        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
@@ -164,6 +168,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, save_name: str = "", show_annotations: bool = False):
         """
         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,16 +28,16 @@ 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
         """
         image = data.image
         predictions = self.__model.predict(
@@ -50,7 +50,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
@@ -84,7 +83,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", 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, save_name, show_annotations)
 
     while not controller.is_exit_requested():
         controller.check_pause()