diff --git a/modules/detect_target/detect_target_contour.py b/modules/detect_target/detect_target_contour.py new file mode 100644 index 00000000..c283a69a --- /dev/null +++ b/modules/detect_target/detect_target_contour.py @@ -0,0 +1,139 @@ +""" +Detects objects using the provided model. +""" + +import time + +import cv2 +import numpy as np + +from . import base_detect_target +from .. import image_and_time +from .. import detections_and_time +from ..common.modules.logger import logger + + +MIN_CONTOUR_AREA = 100 +MAX_CIRCULARITY = 1.3 +MIN_CIRCULARITY = 0.7 + +UPPER_BLUE = np.array([130, 255, 255]) +LOWER_BLUE = np.array([100, 50, 50]) + +CONFIDENCE = 1.0 +LABEL = 0 + + +class DetectTargetContour(base_detect_target.BaseDetectTarget): + """ + Predicts and locates landing pads using the classical computer vision methodology. + """ + + def __init__( + self, image_logger: logger.Logger, show_annotations: bool = False, save_name: str = "" + ) -> None: + """ + image_logger: Log annotated images. + show_annotations: Display annotated images. + save_name: filename prefix for logging detections and annotated images. + """ + self.__counter = 0 + self.__show_annotations = show_annotations + self.__filename_prefix = "" + self.__logger = image_logger + + if save_name != "": + self.__filename_prefix = save_name + "_" + str(int(time.time())) + "_" + + def detect_landing_pads_contours( + self, image_and_time_data: image_and_time.ImageAndTime + ) -> tuple[True, detections_and_time.DetectionsAndTime, np.ndarray] | tuple[False, None, None]: + """ + Detects landing pads using contours/classical CV. + + image_and_time_data: Data for the current image and time. + timestamp: Timestamp for the detections. + + Return: Success, the DetectionsAndTime object, and the annotated image. + """ + image = image_and_time_data.image + timestamp = image_and_time_data.timestamp + + hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + mask = cv2.inRange(hsv_image, LOWER_BLUE, UPPER_BLUE) + + contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + if len(contours) == 0: + return False, None, None + + result, detections = detections_and_time.DetectionsAndTime.create(timestamp) + if not result: + return False, None, None + + image_annotated = image + for i, contour in enumerate(contours): + contour_area = cv2.contourArea(contour) + + if contour_area < MIN_CONTOUR_AREA: + continue + + (x, y), radius = cv2.minEnclosingCircle(contour) + + enclosing_area = np.pi * (radius**2) + circularity = contour_area / enclosing_area + + if circularity < MIN_CIRCULARITY or circularity > MAX_CIRCULARITY: + continue + + x, y, w, h = cv2.boundingRect(contour) + bounds = np.array([x, y, x + w, y + h]) + + # Create a Detection object and append it to detections + result, detection = detections_and_time.Detection.create(bounds, LABEL, CONFIDENCE) + + if not result: + return False, None, None + + detections.append(detection) + + # Annotate the image + cv2.rectangle(image_annotated, (x, y), (x + w, y + h), (0, 0, 255), 2) + cv2.putText( + image_annotated, + f"landing-pad {i+1}", + (x, y - 10), + cv2.FONT_HERSHEY_SIMPLEX, + 0.9, + (0, 0, 255), + 2, + ) + + return True, detections, image_annotated + + def run( + self, data: image_and_time.ImageAndTime + ) -> tuple[True, detections_and_time.DetectionsAndTime] | tuple[False, None]: + """ + Runs object detection on the provided image and returns the detections. + + data: Image with a timestamp. + + Return: Success and the detections. + """ + + result, detections, image_annotated = self.detect_landing_pads_contours(data) + + if not result: + return False, None + + # Logging + if self.__filename_prefix != "": + filename = self.__filename_prefix + str(self.__counter) + self.__logger.save_image(image_annotated, filename) + self.__counter += 1 + + if self.__show_annotations: + cv2.imshow("Annotated", image_annotated) # type: ignore + + return True, detections diff --git a/modules/detect_target/detect_target_factory.py b/modules/detect_target/detect_target_factory.py index 376456cb..b241cdaa 100644 --- a/modules/detect_target/detect_target_factory.py +++ b/modules/detect_target/detect_target_factory.py @@ -6,6 +6,7 @@ from . import base_detect_target from . import detect_target_brightspot +from . import detect_target_contour from . import detect_target_ultralytics from ..common.modules.logger import logger @@ -17,6 +18,7 @@ class DetectTargetOption(enum.Enum): ML_ULTRALYTICS = 0 CV_BRIGHTSPOT = 1 + CV_CONTOUR = 2 def create_detect_target( @@ -47,5 +49,11 @@ def create_detect_target( show_annotations, save_name, ) + case DetectTargetOption.CV_CONTOUR: + return True, detect_target_contour.DetectTargetContour( + local_logger, + show_annotations, + save_name, + ) return False, None diff --git a/tests/unit/generate_detect_target_contour.py b/tests/unit/generate_detect_target_contour.py new file mode 100644 index 00000000..c1677f4a --- /dev/null +++ b/tests/unit/generate_detect_target_contour.py @@ -0,0 +1,203 @@ +""" +Helper functions for `test_detect_target_contour.py`. + +""" + +import cv2 +import math +import numpy as np + +from modules import detections_and_time +from modules import image_and_time + + +LANDING_PAD_COLOUR_BLUE = (100, 50, 50) # BGR + + +class LandingPadImageConfig: + """ + Represents the data required to define and generate a landing pad. + """ + + def __init__( + self, + centre: tuple[int, int], + axis: tuple[int, int], + blur: bool, + angle: float, + ): + """ + centre: The pixel coordinates representing the centre of the landing pad. + axis: The pixel lengths of the semi-major axes of the ellipse. + blur: Indicates whether the landing pad should have a blur effect. + angle: The rotation angle of the landing pad in degrees clockwise, where 0.0 degrees + is where both semi major and minor are aligned with the x and y-axis respectively (0.0 <= angle <= 360.0). + """ + self.centre = centre + self.axis = axis + self.blur = blur + self.angle = angle + + +class NumpyImage: + """ + Holds the numpy array which represents an image. + """ + + def __init__(self, image: np.ndarray): + """ + image: A numpy array that represents the image. + """ + self.image = image + + +class BoundingBox: + """ + Holds the data that define the generated bounding boxes. + """ + + def __init__(self, top_left: tuple[float, float], bottom_right: tuple[float, float]): + """ + top_left: The pixel coordinates representing the top left corner of the bounding box on an image. + bottom_right: pixel coordinates representing the bottom right corner of the bounding box on an image. + """ + self.top_left = top_left + self.bottom_right = bottom_right + + +class InputImageAndTimeAndExpectedBoundingBoxes: + """ + Struct to hold the data needed to perform the tests. + """ + + def __init__(self, image_and_time_data: image_and_time.ImageAndTime, bounding_box_list: list): + """ + image_and_time_data: ImageAndTime object containing the image and timestamp + bounding_box_list: A list that holds expected bounding box coordinates. + Given in the following format: + [conf, label, top_left_x, top_left_y, bottom_right_x, bottom_right_y] + """ + self.image_and_time_data = image_and_time_data + self.bounding_box_list = bounding_box_list + + +def add_blurred_landing_pad( + background: np.ndarray, landing_data: LandingPadImageConfig +) -> NumpyImage: + """ + Blurs an image and adds a singular landing pad to the background. + + background: A numpy image. + landing_data: Landing pad data for the landing pad to be blurred and added. + + Returns: Image with the landing pad. + """ + x, y = background.shape[:2] + + mask = np.zeros((x, y), np.uint8) + mask = cv2.ellipse( + mask, + landing_data.centre, + landing_data.axis, + landing_data.angle, + 0, + 360, + 255, + -1, + ) + + mask = cv2.blur(mask, (25, 25), 7) + + alpha = mask[:, :, np.newaxis] / 255.0 + # Brings the image back to its original colour. + fg = np.full(background.shape, LANDING_PAD_COLOUR_BLUE, dtype=np.uint8) + + blended = (background * (1 - alpha) + fg * alpha).astype(np.uint8) + return NumpyImage(blended) + + +def draw_landing_pad( + image: np.ndarray, landing_data: LandingPadImageConfig +) -> tuple[NumpyImage, BoundingBox]: + """ + Draws a single landing pad on the provided image and saves the bounding box coordinates to a text file. + + image: The image to add a landing pad to. + landing_data: Landing pad data for the landing pad to be added. + + Returns: Image with landing pad and the bounding box for the drawn landing pad. + """ + centre_x, centre_y = landing_data.centre + axis_x, axis_y = landing_data.axis + angle_in_rad = math.radians(landing_data.angle) + + ux = axis_x * math.cos(angle_in_rad) + uy = axis_x * math.sin(angle_in_rad) + + vx = axis_y * math.sin(angle_in_rad) + vy = axis_y * math.cos(angle_in_rad) + + width = 2 * math.sqrt(ux**2 + vx**2) + height = 2 * math.sqrt(uy**2 + vy**2) + + top_left = (max(centre_x - (0.5) * width, 0), max(centre_y - (0.5) * height, 0)) + bottom_right = ( + min(centre_x + (0.5) * width, image.shape[1]), + min(centre_y + (0.5) * height, image.shape[0]), + ) + + bounding_box = BoundingBox(top_left, bottom_right) + + if landing_data.blur: + image = add_blurred_landing_pad(image, landing_data) + return image, bounding_box + + image = cv2.ellipse( + image, + landing_data.centre, + landing_data.axis, + landing_data.angle, + 0, + 360, + LANDING_PAD_COLOUR_BLUE, + -1, + ) + return NumpyImage(image), bounding_box + + +def create_test( + landing_list: list[LandingPadImageConfig], +) -> InputImageAndTimeAndExpectedBoundingBoxes: + """ + Generates test cases given a data set. + + landing_list: List of landing pad data to be generated. + + Returns: The image and expected bounding box. + """ + image = np.full(shape=(1000, 2000, 3), fill_value=255, dtype=np.uint8) + confidence_and_label = [1, 0] + + # List to hold the bounding boxes. + # boxes_list = [confidence, label, top_left_x, top_left_y, bottom_right_x, bottom_right_y] + boxes_list = [] + + for landing_data in landing_list: + image_wrapper, bounding_box = draw_landing_pad(image, landing_data) + image = image_wrapper.image + boxes_list.append( + confidence_and_label + list(bounding_box.top_left + bounding_box.bottom_right) + ) + + # Sorts by the area of the bounding box + boxes_list = sorted( + boxes_list, reverse=True, key=lambda box: abs((box[4] - box[2]) * (box[5] - box[3])) + ) + + image = image.astype(np.uint8) + result, image_and_time_data = image_and_time.ImageAndTime.create(image) + + assert result + assert image_and_time_data is not None + + return InputImageAndTimeAndExpectedBoundingBoxes(image_and_time_data, boxes_list) diff --git a/tests/unit/test_detect_target_contour.py b/tests/unit/test_detect_target_contour.py new file mode 100644 index 00000000..3f0c7d39 --- /dev/null +++ b/tests/unit/test_detect_target_contour.py @@ -0,0 +1,237 @@ +""" +Test contour detection module. +""" + +import numpy as np +import pytest + +from modules import detections_and_time +from modules.detect_target import detect_target_contour +from modules.common.modules.logger import logger +from tests.unit import generate_detect_target_contour + +BOUNDING_BOX_PRECISION_TOLERANCE = -1 # Tolerance > 1 +CONFIDENCE_PRECISION_TOLERANCE = 2 +LOGGER_NAME = "" + + +# Test functions use test fixture signature names and access class privates +# pylint: disable=protected-access,redefined-outer-name, duplicate-code + + +@pytest.fixture +def single_circle() -> generate_detect_target_contour.InputImageAndTimeAndExpectedBoundingBoxes: + """ + Loads the data for the single basic circle. + """ + options = [ + generate_detect_target_contour.LandingPadImageConfig( + centre=(300, 400), axis=(200, 200), blur=False, angle=0 + ) + ] + + test_data = generate_detect_target_contour.create_test(options) + yield test_data + + +@pytest.fixture +def blurry_circle() -> generate_detect_target_contour.InputImageAndTimeAndExpectedBoundingBoxes: + """ + Loads the data for the single blury circle. + """ + options = [ + generate_detect_target_contour.LandingPadImageConfig( + centre=(1000, 500), axis=(423, 423), blur=True, angle=0 + ), + ] + + test_data = generate_detect_target_contour.create_test(options) + yield test_data + + +@pytest.fixture +def stretched_circle() -> generate_detect_target_contour.InputImageAndTimeAndExpectedBoundingBoxes: + """ + Loads the data for the single stretched circle. + """ + options = [ + generate_detect_target_contour.LandingPadImageConfig( + centre=(1000, 500), axis=(383, 405), blur=False, angle=0 + ) + ] + + test_data = generate_detect_target_contour.create_test(options) + yield test_data + + +@pytest.fixture +def multiple_circles() -> generate_detect_target_contour.InputImageAndTimeAndExpectedBoundingBoxes: + """ + Loads the data for the multiple stretched circles. + """ + options = [ + generate_detect_target_contour.LandingPadImageConfig( + centre=(997, 600), axis=(300, 300), blur=False, angle=0 + ), + generate_detect_target_contour.LandingPadImageConfig( + centre=(1590, 341), axis=(250, 250), blur=False, angle=0 + ), + generate_detect_target_contour.LandingPadImageConfig( + centre=(200, 500), axis=(50, 45), blur=True, angle=0 + ), + generate_detect_target_contour.LandingPadImageConfig( + centre=(401, 307), axis=(200, 150), blur=True, angle=0 + ), + ] + + test_data = generate_detect_target_contour.create_test(options) + yield test_data + + +@pytest.fixture() +def detector() -> detect_target_contour.DetectTargetContour: # type: ignore + """ + Construct DetectTargetContour. + """ + result, test_logger = logger.Logger.create(LOGGER_NAME, False) + + assert result + assert test_logger is not None + + detection = detect_target_contour.DetectTargetContour(test_logger, False) + yield detection # type: ignore + + +def compare_detections( + actual_and_expected_detections: generate_detect_target_contour.InputImageAndTimeAndExpectedBoundingBoxes, +) -> None: + """ + Compare expected and actual detections. + + actual_and_expected_detections: Test data containing both actual image and time and expected bounding boxes. + """ + actual = actual_and_expected_detections.image_and_time_data.detector + expected = actual_and_expected_detections.bounding_box_list + + assert len(actual) == len(expected) + + # Ordered for the mapping to the corresponding detections + sorted_actual_detections = sorted( + actual, + reverse=True, + key=lambda box: abs((box.x_2 - box.x_1) * (box.y_2 - box.y_1)), + ) + + for i, expected_detection in enumerate(expected): + actual_detection = sorted_actual_detections[i] + + # Check label and confidence + assert actual_detection.label == expected_detection[1] + np.testing.assert_almost_equal( + actual_detection.confidence, + expected_detection[0], + decimal=CONFIDENCE_PRECISION_TOLERANCE, + ) + + # Check bounding box coordinates + np.testing.assert_almost_equal( + actual_detection.x_1, + expected_detection[2], + decimal=BOUNDING_BOX_PRECISION_TOLERANCE, + ) + + np.testing.assert_almost_equal( + actual_detection.y_1, + expected_detection[3], + decimal=BOUNDING_BOX_PRECISION_TOLERANCE, + ) + + np.testing.assert_almost_equal( + actual_detection.x_2, + expected_detection[4], + decimal=BOUNDING_BOX_PRECISION_TOLERANCE, + ) + + np.testing.assert_almost_equal( + actual_detection.y_2, + expected_detection[5], + decimal=BOUNDING_BOX_PRECISION_TOLERANCE, + ) + + +class TestDetector: + """ + Tests `DetectTarget.run()` . + """ + + def test_single_circle( + self, + detector: detect_target_contour.DetectTargetContour, + single_circle: generate_detect_target_contour.InputImageAndTimeAndExpectedBoundingBoxes, + ) -> None: + """ + Run the detection for the single circular landing pad. + """ + # Run + result, actual = detector.run(single_circle.image_and_time_data) + + # Test + assert result + assert actual is not None + + # Create new object with actual detections + test_data = generate_detect_target_contour.InputImageAndTimeAndExpectedBoundingBoxes( + actual, single_circle.bounding_box_list + ) + compare_detections(test_data) + + def test_blurry_circle( + self, + detector: detect_target_contour.DetectTargetContour, + blurry_circle: generate_detect_target_contour.InputImageAndTimeAndExpectedBoundingBoxes, + ) -> None: + """ + Run the detection for the blury cicular circle. + """ + # Run + result, actual = detector.run(blurry_circle.image_and_time_data) + + # Test + assert result + assert actual is not None + + compare_detections(blurry_circle) + + def test_stretched_circle( + self, + detector: detect_target_contour.DetectTargetContour, + stretched_circle: generate_detect_target_contour.InputImageAndTimeAndExpectedBoundingBoxes, + ) -> None: + """ + Run the detection for a single stretched circular landing pad. + """ + # Run + result, actual = detector.run(stretched_circle.image_and_time_data) + + # Test + assert result + assert actual is not None + + compare_detections(stretched_circle) + + def test_multiple_circles( + self, + detector: detect_target_contour.DetectTargetContour, + multiple_circles: generate_detect_target_contour.InputImageAndTimeAndExpectedBoundingBoxes, + ) -> None: + """ + Run the detection for the multiple landing pads. + """ + # Run + result, actual = detector.run(multiple_circles.image_and_time_data) + + # Test + assert result + assert actual is not None + + compare_detections(multiple_circles.bounding_box_list)