From 72f0ae0667ca5e4ee958aea5e08ab52b1bf440aa Mon Sep 17 00:00:00 2001 From: SkaisteMot <118757021+SkaisteMot@users.noreply.github.com> Date: Sat, 26 Oct 2024 16:22:33 +0100 Subject: [PATCH] Colour detection updates (#1) * removing line * hopefully final colour w linting okay? well see --------- Co-authored-by: Skaiste Motiejunaite --- Algorithms/Objects/colour_detection.py | 75 +++++++++++++++++--------- 1 file changed, 51 insertions(+), 24 deletions(-) diff --git a/Algorithms/Objects/colour_detection.py b/Algorithms/Objects/colour_detection.py index 5947dcf..6023f30 100644 --- a/Algorithms/Objects/colour_detection.py +++ b/Algorithms/Objects/colour_detection.py @@ -4,53 +4,80 @@ Code from: https://www.geeksforgeeks.org/multiple-color-detection-in-real-time-using-python-opencv/ """ -from cv2 import cv2 +import sys +import cv2 import numpy as np import pandas as pd # Load the colour ranges from the CSV file def load_colour_ranges(csv_file): + """Load in CSV file of colours""" colour_data = pd.read_csv(csv_file) - colour_ranges = {} + colours = {} for _, row in colour_data.iterrows(): - colour_ranges[row['colour']] = ( + colours[row['colour']] = ( [row['h_min'], row['s_min'], row['v_min']], [row['h_max'], row['s_max'], row['v_max']] ) - return colour_ranges + return colours # Load colour ranges from CSV file colour_ranges = load_colour_ranges('../../Datasets/colour_ranges.csv') -# Use all colours from the CSV file -selected_colours = list(colour_ranges.keys()) # is this needed? cant we just use the csv since all colours are used regardless? +def create_mask(hsv_img, lower, upper): + """Create a mask for the selected colour""" + return cv2.inRange(hsv_img, lower, upper) + +def draw_bounding_box(image, contour, colour): + """Draw bounding box and label for detected colour""" + x, y, w, h = cv2.boundingRect(contour) + cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 0), 2, lineType=cv2.LINE_AA) + cv2.putText(image, colour, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2) + +def count_colours(image, contours, colour, min_area): + """Count detected colours based on contours""" + colour_count = 0 + for contour in contours: + area = cv2.contourArea(contour) + if area > min_area: + colour_count += 1 + draw_bounding_box(image, contour, colour) + return colour_count + +def display_colour_counts(image, colour_counts): + """Display the counts of detected colours on the image""" + y_offset = 30 # Start drawing from the top of the image + for i, (colour, count) in enumerate(colour_counts.items()): + cv2.putText( + image, + f"{colour}: {count}", + (10, y_offset + i * 20), + cv2.FONT_HERSHEY_SIMPLEX, + 0.6, + (0, 0, 0), + 2 + ) -# Function to detect colours and draw bounding boxes with colour names -def detect_and_draw(image, colour_ranges, selected_colours, min_area=300): - """Detect colours and draw a bounding box""" - # Convert the image to HSV format +def detect_and_draw(image, colours, min_area=300): + """Detect colours and draw bounding boxes with colour names""" hsv_img = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + colour_counts = {} # Dictionary to store the count of each colour detected - for colour in selected_colours: - lower, upper = colour_ranges[colour] + for colour, (lower, upper) in colours.items(): lower = np.array(lower, dtype="uint8") upper = np.array(upper, dtype="uint8") # Create a mask for the selected colour - mask = cv2.inRange(hsv_img, lower, upper) + mask = create_mask(hsv_img, lower, upper) # Find contours (these will be the detected objects) contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - # Draw bounding boxes around detected colours with a black line - for contour in contours: - area = cv2.contourArea(contour) # Calculate the area of the contour + # Count detected colours + colour_counts[colour] = count_colours(image, contours, colour, min_area) - if area > min_area: # Check if the area is above the minimum threshold - x, y, w, h = cv2.boundingRect(contour) - cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 0), 2, lineType=cv2.LINE_AA) - # Display the colour name above the bounding box - cv2.putText(image, colour, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2) + # Display the colour counts at the top left of the image + display_colour_counts(image, colour_counts) return image @@ -59,7 +86,7 @@ def detect_and_draw(image, colour_ranges, selected_colours, min_area=300): if not cap.isOpened(): print("Error: Could not open video stream.") - exit() + sys.exit() # Use sys.exit() instead of exit() try: while True: @@ -70,7 +97,7 @@ def detect_and_draw(image, colour_ranges, selected_colours, min_area=300): break # Detect and draw bounding boxes for selected colours - result_frame = detect_and_draw(frame, colour_ranges, selected_colours, min_area=300) + result_frame = detect_and_draw(frame, colour_ranges, min_area=300) # Show the result cv2.imshow("Detected Colours", result_frame) @@ -78,7 +105,7 @@ def detect_and_draw(image, colour_ranges, selected_colours, min_area=300): # Break the loop on 'q' key press if cv2.waitKey(1) & 0xFF == ord('q'): break -except Exception as e: +except Exception as e: # Catch all exceptions, but be cautious with this approach print(f"Error: {e}") finally: # Release the video capture object and close all OpenCV windows