On limelight python pipeline

[wmbelk]

To look for purple and white

import cv2
import numpy as np

def detect_pole(image, lower_color, upper_color, min_area, max_area):
    # Convert the strip to HSV color space
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    
    # Create a mask for the color range
    mask = cv2.inRange(hsv, lower_color, upper_color)
    
    # Apply Gaussian blur to reduce noise
    blurred = cv2.GaussianBlur(mask, (5, 5), 0)
    
    # Find contours in the mask
    contours, _ = cv2.findContours(blurred, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Initialize variables
    largest_contour = None
    largest_area = 0
    
    # Find the largest contour within the specified area range
    for contour in contours:
        area = cv2.contourArea(contour)
        if min_area < area < max_area and area > largest_area:
            largest_contour = contour
            largest_area = area
    
    return largest_contour

def draw_pole(image, contour, color):
    if contour is not None:
        x, y, w, h = cv2.boundingRect(contour)
        cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
        cv2.putText(image, "Pole Detected", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

def runPipeline(image, llrobot):
    # Initialize variables
    largest_contour = np.array([[]])
    llpython = [0, 0, 0, 0, 0, 0, 0, 0]
    
    try:
        # Define the horizontal strip (middle 20% of the image)
        height, width = image.shape[:2]
        strip_start = int(height * 0.4)
        strip_end = int(height * 0.6)
        image_strip = image[strip_start:strip_end, :]
        
        # Define color ranges in HSV
        lower_purple = np.array([130, 50, 50])
        upper_purple = np.array([160, 255, 255])
        lower_white = np.array([0, 0, 200])
        upper_white = np.array([180, 30, 255])
        
        # Detect purple pole
        purple_contour = detect_pole(image_strip, lower_purple, upper_purple, 100, 5000)
        
        if purple_contour is not None:
            draw_pole(image, purple_contour, (255, 0, 255))  # Magenta color
            largest_contour = purple_contour
        else:
            # If purple pole not found, detect white pole
            white_contour = detect_pole(image_strip, lower_white, upper_white, 200, 10000)
            if white_contour is not None:
                draw_pole(image, white_contour, (255, 255, 255))  # White color
                largest_contour = white_contour
        
        # Draw the horizontal strip boundaries
        cv2.line(image, (0, strip_start), (width, strip_start), (0, 255, 0), 2)
        cv2.line(image, (0, strip_end), (width, strip_end), (0, 255, 0), 2)
        
        # If a pole is detected, update llpython with its properties
        if largest_contour.size > 0:
            x, y, w, h = cv2.boundingRect(largest_contour)
            llpython = [1, x, y, w, h, 0, 0, 0]
        
    except Exception as e:
        # Handle any unexpected errors
        cv2.putText(image, f"Error: {str(e)}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    
    return largest_contour, image, llpython