diff --git a/config.yaml b/config.yaml index 6ca2b687..c6d4f59e 100644 --- a/config.yaml +++ b/config.yaml @@ -22,27 +22,29 @@ video_input: detect_target: worker_count: 1 - option: 0 # 0 is for Ultralytics (from detect_target_factory.py) - device: 0 - model_path: "tests/model_example/yolov8s_ultralytics_pretrained_default.pt" # See autonomy OneDrive for latest model + option: 0 # 0 is for Ultralytics and 1 is for brightspot (from detect_target_factory.py) save_prefix: "log_comp" - -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 f927f220..25f68a72 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 @@ -109,13 +111,26 @@ def main() -> int: 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 + 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"] @@ -244,12 +259,10 @@ def main() -> int: 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_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], diff --git a/modules/detect_target/detect_target_factory.py b/modules/detect_target/detect_target_factory.py index ed292db3..376456cb 100644 --- a/modules/detect_target/detect_target_factory.py +++ b/modules/detect_target/detect_target_factory.py @@ -20,13 +20,14 @@ class DetectTargetOption(enum.Enum): def create_detect_target( + save_name: str, + show_annotations: bool, 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, ) -> tuple[bool, base_detect_target.BaseDetectTarget | None]: """ Construct detect target class at runtime. @@ -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_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/modules/detect_target/detect_target_worker.py b/modules/detect_target/detect_target_worker.py index 8ce5cbc4..d4c2ec05 100644 --- a/modules/detect_target/detect_target_worker.py +++ b/modules/detect_target/detect_target_worker.py @@ -7,17 +7,20 @@ 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, - show_annotations: bool, save_name: str, + show_annotations: bool, + detect_target_option: detect_target_factory.DetectTargetOption, + config: ( + detect_target_brightspot.DetectTargetBrightspotConfig + | detect_target_ultralytics.DetectTargetUltralyticsConfig + ), input_queue: queue_proxy_wrapper.QueueProxyWrapper, output_queue: queue_proxy_wrapper.QueueProxyWrapper, controller: worker_controller.WorkerController, @@ -43,13 +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, - device, - model_path, - override_full, + config, local_logger, - show_annotations, - save_name, ) if not result: 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 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