From 9f21e4b3456218a0f1a376ad5cbe4fb8c14a41cb Mon Sep 17 00:00:00 2001 From: Siddh Patel <47033295+siddhp1@users.noreply.github.com> Date: Tue, 21 Jan 2025 23:02:54 -0500 Subject: [PATCH 1/5] Create ultralytics config class --- .../detect_target_ultralytics.py | 35 +++++++++++++++---- tests/unit/test_detect_target_ultralytics.py | 8 +++-- 2 files changed, 34 insertions(+), 9 deletions(-) diff --git a/modules/detect_target/detect_target_ultralytics.py b/modules/detect_target/detect_target_ultralytics.py index 4f9bddc1..d26e762c 100644 --- a/modules/detect_target/detect_target_ultralytics.py +++ b/modules/detect_target/detect_target_ultralytics.py @@ -13,9 +13,9 @@ from ..common.modules.logger import logger -class DetectTargetUltralytics(base_detect_target.BaseDetectTarget): +class DetectTargetUltralyticsConfig: """ - Contains the YOLOv8 model for prediction. + Configuration for DetectTargetUltralytics. """ def __init__( @@ -23,6 +23,27 @@ def __init__( device: "str | int", model_path: str, override_full: bool, + ) -> None: + """ + Initializes the configuration for DetectTargetUltralytics. + + 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. + """ + self.device = device + self.model_path = model_path + self.override_full = override_full + + +class DetectTargetUltralytics(base_detect_target.BaseDetectTarget): + """ + Contains the YOLOv8 model for prediction. + """ + + def __init__( + self, + config: DetectTargetUltralyticsConfig, local_logger: logger.Logger, show_annotations: bool = False, save_name: str = "", @@ -34,14 +55,14 @@ def __init__( 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.__device = config.device self.__enable_half_precision = not self.__device == "cpu" + self.__model = ultralytics.YOLO(config.model_path) + if config.override_full: + self.__enable_half_precision = False + self.__counter = 0 self.__local_logger = local_logger 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())) + "_" diff --git a/tests/unit/test_detect_target_ultralytics.py b/tests/unit/test_detect_target_ultralytics.py index b7159086..c73d88d2 100644 --- a/tests/unit/test_detect_target_ultralytics.py +++ b/tests/unit/test_detect_target_ultralytics.py @@ -114,9 +114,13 @@ def detector() -> detect_target_ultralytics.DetectTargetUltralytics: # type: ig assert result assert test_logger is not None - detection = detect_target_ultralytics.DetectTargetUltralytics( - DEVICE, str(MODEL_PATH), OVERRIDE_FULL, test_logger + config = detect_target_ultralytics.DetectTargetUltralyticsConfig( + DEVICE, + str(MODEL_PATH), + OVERRIDE_FULL, ) + + detection = detect_target_ultralytics.DetectTargetUltralytics(config, test_logger) yield detection # type: ignore From 67862df1c5b250a33082a68e21c5b411590d5058 Mon Sep 17 00:00:00 2001 From: Siddh Patel <47033295+siddhp1@users.noreply.github.com> Date: Tue, 21 Jan 2025 23:04:00 -0500 Subject: [PATCH 2/5] Add configs to target detection pipeline --- config.yaml | 6 +- main_2024.py | 75 +++++++++++++++++-- .../detect_target/detect_target_factory.py | 12 +-- modules/detect_target/detect_target_worker.py | 13 ++-- 4 files changed, 85 insertions(+), 21 deletions(-) diff --git a/config.yaml b/config.yaml index 6ca2b687..0be5c846 100644 --- a/config.yaml +++ b/config.yaml @@ -22,10 +22,12 @@ video_input: detect_target: worker_count: 1 - option: 0 # 0 is for Ultralytics (from detect_target_factory.py) + option: 0 # 0 is for Ultralytics and 1 is for brightspot (from detect_target_factory.py) + save_prefix: "log_comp" + +detect_ultralytics: device: 0 model_path: "tests/model_example/yolov8s_ultralytics_pretrained_default.pt" # See autonomy OneDrive for latest model - save_prefix: "log_comp" detect_brightspot: brightspot_percentile_threshold: 99.9 diff --git a/main_2024.py b/main_2024.py index f927f220..a45b1529 100644 --- a/main_2024.py +++ b/main_2024.py @@ -16,8 +16,10 @@ from modules.common.modules.camera import camera_opencv from modules.common.modules.camera import camera_picamera2 from modules.communications import communications_worker +from modules.detect_target import detect_target_brightspot from modules.detect_target import detect_target_factory from modules.detect_target import detect_target_worker +from modules.detect_target import detect_target_ultralytics from modules.flight_interface import flight_interface_worker from modules.video_input import video_input_worker from modules.data_merge import data_merge_worker @@ -105,17 +107,48 @@ def main() -> int: config["video_input"]["image_name"] if config["video_input"]["log_images"] else None ) - DETECT_TARGET_WORKER_COUNT = config["detect_target"]["worker_count"] DETECT_TARGET_OPTION = detect_target_factory.DetectTargetOption( config["detect_target"]["option"] ) - DETECT_TARGET_DEVICE = "cpu" if args.cpu else config["detect_target"]["device"] - DETECT_TARGET_MODEL_PATH = config["detect_target"]["model_path"] - DETECT_TARGET_OVERRIDE_FULL_PRECISION = args.full DETECT_TARGET_SAVE_PREFIX = str( pathlib.Path(logging_path, config["detect_target"]["save_prefix"]) ) DETECT_TARGET_SHOW_ANNOTATED = args.show_annotated + DETECT_TARGET_WORKER_COUNT = config["detect_target"]["worker_count"] + + DETECT_TARGET_ULTRALYTICS_DEVICE = ( + "cpu" if args.cpu else config["detect_ultralytics"]["device"] + ) + DETECT_TARGET_ULTRALYTICS_MODEL_PATH = config["detect_ultralytics"]["model_path"] + DETECT_TARGET_ULTRALYTICS_OVERRIDE_FULL_PRECISION = args.full + + DETECT_TARGET_BRIGHTSPOT_PERCENTILE_THRESHOLD = config["detect_brightspot"][ + "brightspot_percentile_threshold" + ] + DETECT_TARGET_BRIGHTSPOT_FILTER_BY_COLOR = config["detect_brightspot"]["filter_by_color"] + DETECT_TARGET_BRIGHTSPOT_BLOB_COLOR = config["detect_brightspot"]["blob_color"] + DETECT_TARGET_BRIGHTSPOT_FILTER_BY_CIRCULARITY = config["detect_brightspot"][ + "filter_by_circularity" + ] + DETECT_TARGET_BRIGHTSPOT_MIN_CIRCULARITY = config["detect_brightspot"]["min_circularity"] + DETECT_TARGET_BRIGHTSPOT_MAX_CIRCULARITY = config["detect_brightspot"]["max_circularity"] + DETECT_TARGET_BRIGHTSPOT_FILTER_BY_INERTIA = config["detect_brightspot"][ + "filter_by_inertia" + ] + DETECT_TARGET_BRIGHTSPOT_MIN_INERTIA_RATIO = config["detect_brightspot"][ + "min_inertia_ratio" + ] + DETECT_TARGET_BRIGHTSPOT_MAX_INERTIA_RATIO = config["detect_brightspot"][ + "max_inertia_ratio" + ] + DETECT_TARGET_BRIGHTSPOT_FILTER_BY_CONVEXITY = config["detect_brightspot"][ + "filter_by_convexity" + ] + DETECT_TARGET_BRIGHTSPOT_MIN_CONVEXITY = config["detect_brightspot"]["min_convexity"] + DETECT_TARGET_BRIGHTSPOT_MAX_CONVEXITY = config["detect_brightspot"]["max_convexity"] + DETECT_TARGET_BRIGHTSPOT_FILTER_BY_AREA = config["detect_brightspot"]["filter_by_area"] + DETECT_TARGET_BRIGHTSPOT_MIN_AREA_PIXELS = config["detect_brightspot"]["min_area_pixels"] + DETECT_TARGET_BRIGHTSPOT_MAX_AREA_PIXELS = config["detect_brightspot"]["max_area_pixels"] FLIGHT_INTERFACE_ADDRESS = config["flight_interface"]["address"] FLIGHT_INTERFACE_TIMEOUT = config["flight_interface"]["timeout"] @@ -240,14 +273,42 @@ def main() -> int: # Get Pylance to stop complaining assert video_input_worker_properties is not None + detect_target_brightspot_config = detect_target_brightspot.DetectTargetBrightspotConfig( + DETECT_TARGET_BRIGHTSPOT_PERCENTILE_THRESHOLD, + DETECT_TARGET_BRIGHTSPOT_FILTER_BY_COLOR, + DETECT_TARGET_BRIGHTSPOT_BLOB_COLOR, + DETECT_TARGET_BRIGHTSPOT_FILTER_BY_CIRCULARITY, + DETECT_TARGET_BRIGHTSPOT_MIN_CIRCULARITY, + DETECT_TARGET_BRIGHTSPOT_MAX_CIRCULARITY, + DETECT_TARGET_BRIGHTSPOT_FILTER_BY_INERTIA, + DETECT_TARGET_BRIGHTSPOT_MIN_INERTIA_RATIO, + DETECT_TARGET_BRIGHTSPOT_MAX_INERTIA_RATIO, + DETECT_TARGET_BRIGHTSPOT_FILTER_BY_CONVEXITY, + DETECT_TARGET_BRIGHTSPOT_MIN_CONVEXITY, + DETECT_TARGET_BRIGHTSPOT_MAX_CONVEXITY, + DETECT_TARGET_BRIGHTSPOT_FILTER_BY_AREA, + DETECT_TARGET_BRIGHTSPOT_MIN_AREA_PIXELS, + DETECT_TARGET_BRIGHTSPOT_MAX_AREA_PIXELS, + ) + + detect_target_ultralytics_config = detect_target_ultralytics.DetectTargetUltralyticsConfig( + DETECT_TARGET_ULTRALYTICS_DEVICE, + DETECT_TARGET_ULTRALYTICS_MODEL_PATH, + DETECT_TARGET_ULTRALYTICS_OVERRIDE_FULL_PRECISION, + ) + + match DETECT_TARGET_OPTION: + case detect_target_factory.DetectTargetOption.ML_ULTRALYTICS: + detect_target_config = detect_target_ultralytics_config + case detect_target_factory.DetectTargetOption.CV_BRIGHTSPOT: + detect_target_config = detect_target_brightspot_config + result, detect_target_worker_properties = worker_manager.WorkerProperties.create( count=DETECT_TARGET_WORKER_COUNT, target=detect_target_worker.detect_target_worker, work_arguments=( DETECT_TARGET_OPTION, - DETECT_TARGET_DEVICE, - DETECT_TARGET_MODEL_PATH, - DETECT_TARGET_OVERRIDE_FULL_PRECISION, + detect_target_config, DETECT_TARGET_SHOW_ANNOTATED, DETECT_TARGET_SAVE_PREFIX, ), diff --git a/modules/detect_target/detect_target_factory.py b/modules/detect_target/detect_target_factory.py index ed292db3..e2c23a61 100644 --- a/modules/detect_target/detect_target_factory.py +++ b/modules/detect_target/detect_target_factory.py @@ -21,9 +21,10 @@ class DetectTargetOption(enum.Enum): def create_detect_target( detect_target_option: DetectTargetOption, - device: "str | int", - model_path: str, - override_full: bool, + config: ( + detect_target_brightspot.DetectTargetBrightspotConfig + | detect_target_ultralytics.DetectTargetUltralyticsConfig + ), local_logger: logger.Logger, show_annotations: bool, save_name: str, @@ -34,15 +35,14 @@ def create_detect_target( match detect_target_option: case DetectTargetOption.ML_ULTRALYTICS: return True, detect_target_ultralytics.DetectTargetUltralytics( - device, - model_path, - override_full, + config, local_logger, show_annotations, save_name, ) case DetectTargetOption.CV_BRIGHTSPOT: return True, detect_target_brightspot.DetectTargetBrightspot( + config, local_logger, show_annotations, save_name, diff --git a/modules/detect_target/detect_target_worker.py b/modules/detect_target/detect_target_worker.py index 8ce5cbc4..ffffd7ec 100644 --- a/modules/detect_target/detect_target_worker.py +++ b/modules/detect_target/detect_target_worker.py @@ -7,15 +7,18 @@ from utilities.workers import queue_proxy_wrapper from utilities.workers import worker_controller +from . import detect_target_brightspot from . import detect_target_factory +from . import detect_target_ultralytics from ..common.modules.logger import logger def detect_target_worker( detect_target_option: detect_target_factory.DetectTargetOption, - device: "str | int", - model_path: str, - override_full: bool, + config: ( + detect_target_brightspot.DetectTargetBrightspotConfig + | detect_target_ultralytics.DetectTargetUltralyticsConfig + ), show_annotations: bool, save_name: str, input_queue: queue_proxy_wrapper.QueueProxyWrapper, @@ -44,9 +47,7 @@ def detect_target_worker( result, detector = detect_target_factory.create_detect_target( detect_target_option, - device, - model_path, - override_full, + config, local_logger, show_annotations, save_name, From 0758c96ae66388bda2a8846509169f82662148a7 Mon Sep 17 00:00:00 2001 From: Siddh Patel <47033295+siddhp1@users.noreply.github.com> Date: Thu, 23 Jan 2025 18:34:44 -0500 Subject: [PATCH 3/5] Organize YAML and improve parsing --- config.yaml | 42 ++++++++++++------------- main_2024.py | 88 ++++++++++++---------------------------------------- 2 files changed, 41 insertions(+), 89 deletions(-) diff --git a/config.yaml b/config.yaml index 0be5c846..c6d4f59e 100644 --- a/config.yaml +++ b/config.yaml @@ -24,27 +24,27 @@ detect_target: worker_count: 1 option: 0 # 0 is for Ultralytics and 1 is for brightspot (from detect_target_factory.py) save_prefix: "log_comp" - -detect_ultralytics: - device: 0 - model_path: "tests/model_example/yolov8s_ultralytics_pretrained_default.pt" # See autonomy OneDrive for latest model - -detect_brightspot: - brightspot_percentile_threshold: 99.9 - filter_by_color: True - blob_color: 255 - filter_by_circularity: False - min_circularity: 0.01 - max_circularity: 1 - filter_by_inertia: True - min_inertia_ratio: 0.2 - max_inertia_ratio: 1 - filter_by_convexity: False - min_convexity: 0.01 - max_convexity: 1 - filter_by_area: True - min_area_pixels: 50 - max_area_pixels: 640 + # Ultralytics config (enum 0) + config: + device: 0 + model_path: "tests/model_example/yolov8s_ultralytics_pretrained_default.pt" # See autonomy OneDrive for latest model + # Brightspot config (enum 1) + # config: + # brightspot_percentile_threshold: 99.9 + # filter_by_color: True + # blob_color: 255 + # filter_by_circularity: False + # min_circularity: 0.01 + # max_circularity: 1 + # filter_by_inertia: True + # min_inertia_ratio: 0.2 + # max_inertia_ratio: 1 + # filter_by_convexity: False + # min_convexity: 0.01 + # max_convexity: 1 + # filter_by_area: True + # min_area_pixels: 50 + # max_area_pixels: 640 flight_interface: # Port 5762 connects directly to the simulated auto pilot, which is more realistic diff --git a/main_2024.py b/main_2024.py index a45b1529..25f68a72 100644 --- a/main_2024.py +++ b/main_2024.py @@ -107,6 +107,7 @@ def main() -> int: config["video_input"]["image_name"] if config["video_input"]["log_images"] else None ) + DETECT_TARGET_WORKER_COUNT = config["detect_target"]["worker_count"] DETECT_TARGET_OPTION = detect_target_factory.DetectTargetOption( config["detect_target"]["option"] ) @@ -114,41 +115,22 @@ def main() -> int: pathlib.Path(logging_path, config["detect_target"]["save_prefix"]) ) DETECT_TARGET_SHOW_ANNOTATED = args.show_annotated - DETECT_TARGET_WORKER_COUNT = config["detect_target"]["worker_count"] - - DETECT_TARGET_ULTRALYTICS_DEVICE = ( - "cpu" if args.cpu else config["detect_ultralytics"]["device"] - ) - DETECT_TARGET_ULTRALYTICS_MODEL_PATH = config["detect_ultralytics"]["model_path"] - DETECT_TARGET_ULTRALYTICS_OVERRIDE_FULL_PRECISION = args.full - - DETECT_TARGET_BRIGHTSPOT_PERCENTILE_THRESHOLD = config["detect_brightspot"][ - "brightspot_percentile_threshold" - ] - DETECT_TARGET_BRIGHTSPOT_FILTER_BY_COLOR = config["detect_brightspot"]["filter_by_color"] - DETECT_TARGET_BRIGHTSPOT_BLOB_COLOR = config["detect_brightspot"]["blob_color"] - DETECT_TARGET_BRIGHTSPOT_FILTER_BY_CIRCULARITY = config["detect_brightspot"][ - "filter_by_circularity" - ] - DETECT_TARGET_BRIGHTSPOT_MIN_CIRCULARITY = config["detect_brightspot"]["min_circularity"] - DETECT_TARGET_BRIGHTSPOT_MAX_CIRCULARITY = config["detect_brightspot"]["max_circularity"] - DETECT_TARGET_BRIGHTSPOT_FILTER_BY_INERTIA = config["detect_brightspot"][ - "filter_by_inertia" - ] - DETECT_TARGET_BRIGHTSPOT_MIN_INERTIA_RATIO = config["detect_brightspot"][ - "min_inertia_ratio" - ] - DETECT_TARGET_BRIGHTSPOT_MAX_INERTIA_RATIO = config["detect_brightspot"][ - "max_inertia_ratio" - ] - DETECT_TARGET_BRIGHTSPOT_FILTER_BY_CONVEXITY = config["detect_brightspot"][ - "filter_by_convexity" - ] - DETECT_TARGET_BRIGHTSPOT_MIN_CONVEXITY = config["detect_brightspot"]["min_convexity"] - DETECT_TARGET_BRIGHTSPOT_MAX_CONVEXITY = config["detect_brightspot"]["max_convexity"] - DETECT_TARGET_BRIGHTSPOT_FILTER_BY_AREA = config["detect_brightspot"]["filter_by_area"] - DETECT_TARGET_BRIGHTSPOT_MIN_AREA_PIXELS = config["detect_brightspot"]["min_area_pixels"] - DETECT_TARGET_BRIGHTSPOT_MAX_AREA_PIXELS = config["detect_brightspot"]["max_area_pixels"] + match DETECT_TARGET_OPTION: + case detect_target_factory.DetectTargetOption.ML_ULTRALYTICS: + DETECT_TARGET_CONFIG = detect_target_ultralytics.DetectTargetUltralyticsConfig( + config["detect_target"]["config"]["device"], + config["detect_target"]["config"]["model_path"], + args.full, + ) + case detect_target_factory.DetectTargetOption.CV_BRIGHTSPOT: + DETECT_TARGET_CONFIG = detect_target_brightspot.DetectTargetBrightspotConfig( + **config["detect_target"]["config"] + ) + case _: + main.logger.error( + f"Inputted an invalid detect target option: {DETECT_TARGET_OPTION}", True + ) + return -1 FLIGHT_INTERFACE_ADDRESS = config["flight_interface"]["address"] FLIGHT_INTERFACE_TIMEOUT = config["flight_interface"]["timeout"] @@ -273,44 +255,14 @@ def main() -> int: # Get Pylance to stop complaining assert video_input_worker_properties is not None - detect_target_brightspot_config = detect_target_brightspot.DetectTargetBrightspotConfig( - DETECT_TARGET_BRIGHTSPOT_PERCENTILE_THRESHOLD, - DETECT_TARGET_BRIGHTSPOT_FILTER_BY_COLOR, - DETECT_TARGET_BRIGHTSPOT_BLOB_COLOR, - DETECT_TARGET_BRIGHTSPOT_FILTER_BY_CIRCULARITY, - DETECT_TARGET_BRIGHTSPOT_MIN_CIRCULARITY, - DETECT_TARGET_BRIGHTSPOT_MAX_CIRCULARITY, - DETECT_TARGET_BRIGHTSPOT_FILTER_BY_INERTIA, - DETECT_TARGET_BRIGHTSPOT_MIN_INERTIA_RATIO, - DETECT_TARGET_BRIGHTSPOT_MAX_INERTIA_RATIO, - DETECT_TARGET_BRIGHTSPOT_FILTER_BY_CONVEXITY, - DETECT_TARGET_BRIGHTSPOT_MIN_CONVEXITY, - DETECT_TARGET_BRIGHTSPOT_MAX_CONVEXITY, - DETECT_TARGET_BRIGHTSPOT_FILTER_BY_AREA, - DETECT_TARGET_BRIGHTSPOT_MIN_AREA_PIXELS, - DETECT_TARGET_BRIGHTSPOT_MAX_AREA_PIXELS, - ) - - detect_target_ultralytics_config = detect_target_ultralytics.DetectTargetUltralyticsConfig( - DETECT_TARGET_ULTRALYTICS_DEVICE, - DETECT_TARGET_ULTRALYTICS_MODEL_PATH, - DETECT_TARGET_ULTRALYTICS_OVERRIDE_FULL_PRECISION, - ) - - match DETECT_TARGET_OPTION: - case detect_target_factory.DetectTargetOption.ML_ULTRALYTICS: - detect_target_config = detect_target_ultralytics_config - case detect_target_factory.DetectTargetOption.CV_BRIGHTSPOT: - detect_target_config = detect_target_brightspot_config - result, detect_target_worker_properties = worker_manager.WorkerProperties.create( count=DETECT_TARGET_WORKER_COUNT, target=detect_target_worker.detect_target_worker, work_arguments=( - DETECT_TARGET_OPTION, - detect_target_config, - DETECT_TARGET_SHOW_ANNOTATED, DETECT_TARGET_SAVE_PREFIX, + DETECT_TARGET_SHOW_ANNOTATED, + DETECT_TARGET_OPTION, + DETECT_TARGET_CONFIG, ), input_queues=[video_input_to_detect_target_queue], output_queues=[detect_target_to_data_merge_queue], From 87934d2b9ee730a603ae3017e74a7fde8d83b5a9 Mon Sep 17 00:00:00 2001 From: Siddh Patel <47033295+siddhp1@users.noreply.github.com> Date: Thu, 23 Jan 2025 18:35:06 -0500 Subject: [PATCH 4/5] Reorder args --- modules/detect_target/detect_target_factory.py | 4 ++-- modules/detect_target/detect_target_worker.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/modules/detect_target/detect_target_factory.py b/modules/detect_target/detect_target_factory.py index e2c23a61..376456cb 100644 --- a/modules/detect_target/detect_target_factory.py +++ b/modules/detect_target/detect_target_factory.py @@ -20,14 +20,14 @@ class DetectTargetOption(enum.Enum): def create_detect_target( + save_name: str, + show_annotations: bool, detect_target_option: DetectTargetOption, config: ( detect_target_brightspot.DetectTargetBrightspotConfig | detect_target_ultralytics.DetectTargetUltralyticsConfig ), local_logger: logger.Logger, - show_annotations: bool, - save_name: str, ) -> tuple[bool, base_detect_target.BaseDetectTarget | None]: """ Construct detect target class at runtime. diff --git a/modules/detect_target/detect_target_worker.py b/modules/detect_target/detect_target_worker.py index ffffd7ec..d4c2ec05 100644 --- a/modules/detect_target/detect_target_worker.py +++ b/modules/detect_target/detect_target_worker.py @@ -14,13 +14,13 @@ def detect_target_worker( + save_name: str, + show_annotations: bool, detect_target_option: detect_target_factory.DetectTargetOption, config: ( detect_target_brightspot.DetectTargetBrightspotConfig | detect_target_ultralytics.DetectTargetUltralyticsConfig ), - show_annotations: bool, - save_name: str, input_queue: queue_proxy_wrapper.QueueProxyWrapper, output_queue: queue_proxy_wrapper.QueueProxyWrapper, controller: worker_controller.WorkerController, @@ -46,11 +46,11 @@ def detect_target_worker( local_logger.info("Logger initialized", True) result, detector = detect_target_factory.create_detect_target( + save_name, + show_annotations, detect_target_option, config, local_logger, - show_annotations, - save_name, ) if not result: From d7f4d8dbd3b528cb46c9d336b4e763d070445a63 Mon Sep 17 00:00:00 2001 From: Siddh Patel <47033295+siddhp1@users.noreply.github.com> Date: Thu, 23 Jan 2025 19:06:06 -0500 Subject: [PATCH 5/5] Update integration test --- .../integration/test_detect_target_worker.py | 114 +++++++++++++----- tests/unit/test_detect_target_brightspot.py | 4 +- 2 files changed, 88 insertions(+), 30 deletions(-) diff --git a/tests/integration/test_detect_target_worker.py b/tests/integration/test_detect_target_worker.py index 3a56d922..58858427 100644 --- a/tests/integration/test_detect_target_worker.py +++ b/tests/integration/test_detect_target_worker.py @@ -9,34 +9,70 @@ import cv2 import torch +from modules.detect_target import detect_target_brightspot from modules.detect_target import detect_target_factory from modules.detect_target import detect_target_worker +from modules.detect_target import detect_target_ultralytics from modules import image_and_time from modules import detections_and_time from utilities.workers import queue_proxy_wrapper from utilities.workers import worker_controller -TEST_PATH = pathlib.Path("tests", "model_example") -IMAGE_BUS_PATH = pathlib.Path(TEST_PATH, "bus.jpg") -IMAGE_ZIDANE_PATH = pathlib.Path(TEST_PATH, "zidane.jpg") +BRIGHTSPOT_TEST_PATH = pathlib.Path("tests", "brightspot_example") +IMAGE_BRIGHTSPOT_0_PATH = pathlib.Path(BRIGHTSPOT_TEST_PATH, "ir_detections_0.png") +IMAGE_BRIGHTSPOT_1_PATH = pathlib.Path(BRIGHTSPOT_TEST_PATH, "ir_detections_1.png") + +BRIGHTSPOT_OPTION = detect_target_factory.DetectTargetOption.CV_BRIGHTSPOT +# Logging is identical to detect_target_ultralytics.py +# pylint: disable=duplicate-code +BRIGHTSPOT_CONFIG = detect_target_brightspot.DetectTargetBrightspotConfig( + brightspot_percentile_threshold=99.9, + filter_by_color=True, + blob_color=255, + filter_by_circularity=False, + min_circularity=0.01, + max_circularity=1, + filter_by_inertia=True, + min_inertia_ratio=0.2, + max_inertia_ratio=1, + filter_by_convexity=False, + min_convexity=0.01, + max_convexity=1, + filter_by_area=True, + min_area_pixels=50, + max_area_pixels=640, +) +# pylint: enable=duplicate-code + +ULTRALYTICS_TEST_PATH = pathlib.Path("tests", "model_example") +IMAGE_BUS_PATH = pathlib.Path(ULTRALYTICS_TEST_PATH, "bus.jpg") +IMAGE_ZIDANE_PATH = pathlib.Path(ULTRALYTICS_TEST_PATH, "zidane.jpg") + +ULTRALYTICS_OPTION = detect_target_factory.DetectTargetOption.ML_ULTRALYTICS +ULTRALYTICS_CONFIG = detect_target_ultralytics.DetectTargetUltralyticsConfig( + 0 if torch.cuda.is_available() else "cpu", + pathlib.Path(ULTRALYTICS_TEST_PATH, "yolov8s_ultralytics_pretrained_default.pt"), + False, +) -WORK_COUNT = 3 -DELAY_FOR_SIMULATED_WORKERS = 1 # seconds -DELAY_FOR_CUDA_WARMUP = 20 # seconds - -OPTION = detect_target_factory.DetectTargetOption.ML_ULTRALYTICS -MODEL_PATH = pathlib.Path(TEST_PATH, "yolov8s_ultralytics_pretrained_default.pt") -OVERRIDE_FULL = False SHOW_ANNOTATIONS = False SAVE_NAME = "" # No need to save images +WORK_COUNT = 3 +DELAY_FOR_SIMULATED_ULTRALYTICS_WORKER = 1 # seconds +DELAY_FOR_SIMULATED_BRIGHTSPOT_WORKER = 3 # seconds +DELAY_FOR_CUDA_WARMUP = 20 # seconds + def simulate_previous_worker( image_path: pathlib.Path, in_queue: queue_proxy_wrapper.QueueProxyWrapper ) -> None: """ Place the image into the queue. + + image_path: Path to the image being added to the queue. + in_queue: Input queue. """ image = cv2.imread(str(image_path)) # type: ignore result, value = image_and_time.ImageAndTime.create(image) @@ -45,14 +81,20 @@ def simulate_previous_worker( in_queue.queue.put(value) -def main() -> int: +def run_worker( + option: detect_target_factory.DetectTargetOption, + config: ( + detect_target_brightspot.DetectTargetBrightspotConfig + | detect_target_ultralytics.DetectTargetUltralyticsConfig + ), +) -> None: """ - Main function. + Tests target detection. + + option: Brightspot or Ultralytics. + config: Configuration for respective target detection module. """ # Setup - # Not a constant - # pylint: disable-next=invalid-name - device = 0 if torch.cuda.is_available() else "cpu" controller = worker_controller.WorkerController() mp_manager = mp.Manager() @@ -63,33 +105,39 @@ def main() -> int: worker = mp.Process( target=detect_target_worker.detect_target_worker, args=( - OPTION, - device, - MODEL_PATH, - OVERRIDE_FULL, - SHOW_ANNOTATIONS, SAVE_NAME, + SHOW_ANNOTATIONS, + option, + config, image_in_queue, image_out_queue, controller, ), ) - # Run print("Starting worker") worker.start() - for _ in range(0, WORK_COUNT): - simulate_previous_worker(IMAGE_BUS_PATH, image_in_queue) + if option == ULTRALYTICS_OPTION: + for _ in range(0, WORK_COUNT): + simulate_previous_worker(IMAGE_BUS_PATH, image_in_queue) + + time.sleep(DELAY_FOR_SIMULATED_ULTRALYTICS_WORKER) - time.sleep(DELAY_FOR_SIMULATED_WORKERS) + for _ in range(0, WORK_COUNT): + simulate_previous_worker(IMAGE_ZIDANE_PATH, image_in_queue) - for _ in range(0, WORK_COUNT): - simulate_previous_worker(IMAGE_ZIDANE_PATH, image_in_queue) + # Takes some time for CUDA to warm up + print("Waiting for CUDA to warm up") + time.sleep(DELAY_FOR_CUDA_WARMUP) + elif option == BRIGHTSPOT_OPTION: + for _ in range(0, WORK_COUNT): + simulate_previous_worker(IMAGE_BRIGHTSPOT_0_PATH, image_in_queue) - # Takes some time for CUDA to warm up - print("Waiting for CUDA to warm up") - time.sleep(DELAY_FOR_CUDA_WARMUP) + time.sleep(DELAY_FOR_SIMULATED_BRIGHTSPOT_WORKER * 2.5) + + for _ in range(0, WORK_COUNT): + simulate_previous_worker(IMAGE_BRIGHTSPOT_1_PATH, image_in_queue) controller.request_exit() print("Requested exit") @@ -106,6 +154,14 @@ def main() -> int: image_in_queue.fill_and_drain_queue() worker.join() + +def main() -> int: + """ + Main function. + """ + run_worker(BRIGHTSPOT_OPTION, BRIGHTSPOT_CONFIG) + run_worker(ULTRALYTICS_OPTION, ULTRALYTICS_CONFIG) + return 0 diff --git a/tests/unit/test_detect_target_brightspot.py b/tests/unit/test_detect_target_brightspot.py index 6e0ce99d..2bce6285 100644 --- a/tests/unit/test_detect_target_brightspot.py +++ b/tests/unit/test_detect_target_brightspot.py @@ -35,6 +35,8 @@ BOUNDING_BOX_PRECISION_TOLERANCE = 3 CONFIDENCE_PRECISION_TOLERANCE = 6 +# Config is identical to test_detect_target_worker.py +# pylint: disable=duplicate-code DETECT_TARGET_BRIGHTSPOT_CONFIG = detect_target_brightspot.DetectTargetBrightspotConfig( brightspot_percentile_threshold=99.9, filter_by_color=True, @@ -52,7 +54,7 @@ min_area_pixels=50, max_area_pixels=640, ) - +# pylint: enable=duplicate-code # Test functions use test fixture signature names and access class privates # No enable