# Vehicle Detection and State Estimation using Color-Based Contour Detection """ Vehicle Detection Model using Color-Based Contour Detection This module provides functionality to detect vehicles in Bird's-Eye View (BEV) images by isolating their specific colors (green and blue) and analyzing the shapes (contours) of the colored areas. The detected states can be exported to CSV files. Required Libraries: - opencv-python: For image processing, color segmentation, and contour analysis - numpy: For numerical operations You can install them using pip: pip install opencv-python-headless numpy Example usage: from models.contour_detection_model import ContourDetectionModel model = ContourDetectionModel() annotated_image, vehicle_states = model.detect_vehicles('path/to/image.jpg') model.save_states_to_csv(vehicle_states, 'output.csv') """ import cv2 import numpy as np import math import csv import torch from typing import List, Dict, Tuple, Optional class ContourDetectionModel: """ A vehicle detection model using color-based contour detection. This model detects vehicles by: 1. Converting images to HSV color space 2. Creating color masks for green (ego) and blue (other) vehicles 3. Finding contours in the masked regions 4. Estimating vehicle position and heading from contour geometry """ def __init__(self, green_hsv_range: Tuple[List[int], List[int]] = None, blue_hsv_range: Tuple[List[int], List[int]] = None, min_contour_area: int = 50): """ Initialize the contour detection model. Args: green_hsv_range: Tuple of (lower, upper) HSV ranges for green vehicles blue_hsv_range: Tuple of (lower, upper) HSV ranges for blue vehicles min_contour_area: Minimum contour area to filter out noise """ # Default HSV ranges optimized for highway-env vehicles if green_hsv_range is None: self.lower_green = np.array([50, 100, 100]) self.upper_green = np.array([70, 255, 255]) else: self.lower_green = np.array(green_hsv_range[0]) self.upper_green = np.array(green_hsv_range[1]) if blue_hsv_range is None: # Expanded range for better detection in generated images self.lower_blue = np.array([80, 80, 80]) # More tolerant lower bounds self.upper_blue = np.array([115, 255, 255]) # Wider hue range else: self.lower_blue = np.array(blue_hsv_range[0]) self.upper_blue = np.array(blue_hsv_range[1]) self.min_contour_area = min_contour_area def detect_vehicles(self, image_path: str) -> Tuple[Optional[np.ndarray], List[Dict]]: """ Detect vehicles in an image and estimate their states. Args: image_path (str): Path to the input image Returns: Tuple containing: - annotated_image: Image with detection annotations (or None if error) - vehicle_states: List of dictionaries containing vehicle state information """ return self.estimate_vehicle_states_by_color(image_path) def estimate_vehicle_states_by_color(self, image_path: str) -> Tuple[Optional[np.ndarray], List[Dict]]: """ Detects vehicles in an image based on color, and estimates their position and heading. Args: image_path (str): The path to the input image. Returns: tuple: A tuple containing the annotated image and a list of vehicle states. """ # Load the image try: img = cv2.imread(image_path) if img is None: print(f"Error: Could not read image from path: {image_path}") return None, [] # Create a copy for drawing annotations annotated_img = img.copy() except Exception as e: print(f"Error loading image: {e}") return None, [] # Use the existing detection pipeline return self.detect_from_image(img) def save_states_to_csv(self, states: List[Dict], csv_file_path: str) -> None: """ Saves the list of vehicle states to a CSV file. Args: states (list): A list of dictionaries, where each dictionary is a vehicle's state. csv_file_path (str): The path to the output CSV file. """ # Define the fieldnames for the CSV header. We exclude the bounding box points. fieldnames = ['class', 'position_x', 'position_y', 'heading', 'speed'] try: with open(csv_file_path, 'w', newline='') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=fieldnames, extrasaction='ignore') writer.writeheader() # Write the header row writer.writerows(states) # Write all the state data print(f"\nVehicle states successfully saved to: {csv_file_path}") except IOError as e: print(f"Error writing to CSV file: {e}") def process_image(self, input_image_path: str, output_image_path: str = None, output_csv_path: str = None, verbose: bool = True) -> Tuple[Optional[np.ndarray], List[Dict]]: """ Complete processing pipeline for a single image. Args: input_image_path (str): Path to input image output_image_path (str, optional): Path to save annotated image output_csv_path (str, optional): Path to save CSV file verbose (bool): Whether to print detection results Returns: Tuple of (annotated_image, vehicle_states) """ # Process the image to get states and the annotated image annotated_image, states = self.detect_vehicles(input_image_path) if annotated_image is not None and states: if verbose: print("--- Detected Vehicle States (Contour Method) ---") # Sort states by x-position for consistent ordering states.sort(key=lambda v: v['position_x']) for i, state in enumerate(states): print(f"\nVehicle #{i+1}:") print(f" Class: {state['class']}") print(f" Position (x, y): ({state['position_x']:.2f}, {state['position_y']:.2f})") print(f" Heading (degrees): {state['heading']:.2f}") print(f" Speed: {state['speed']:.2f} (Note: Placeholder value)") # Save the annotated image if path provided if output_image_path: cv2.imwrite(output_image_path, annotated_image) if verbose: print(f"\nAnnotated image saved to: {output_image_path}") # Save the states to a CSV file if path provided if output_csv_path: self.save_states_to_csv(states, output_csv_path) elif not states and verbose: print("No vehicles were detected in the image.") return annotated_image, states def detect_from_image(self, img) -> Tuple[Optional[np.ndarray], List[Dict]]: """ Internal method to detect vehicles from numpy image array. Args: img: Input image in BGR format annotated_img: Copy of image for annotations Returns: Tuple of (annotated_image, vehicle_states) """ annotated_img = img.copy() hsv_img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) # 3. Define color ranges for the vehicles # These values are tuned for the specific green and blue in the provided image. # Format: [Hue, Saturation, Value] # Green vehicle (Ego) lower_green = np.array([50, 100, 100]) upper_green = np.array([70, 255, 255]) # Blue vehicles (Corrected Range) lower_blue = np.array([85, 100, 100]) upper_blue = np.array([110, 255, 255]) # 4. Create masks for each color mask_green = cv2.inRange(hsv_img, lower_green, upper_green) mask_blue = cv2.inRange(hsv_img, lower_blue, upper_blue) # 5. Find contours for each mask separately for robust classification contours_green, _ = cv2.findContours(mask_green, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) contours_blue, _ = cv2.findContours(mask_blue, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # Combine contours with their respective class names all_contours = [] for c in contours_green: all_contours.append((c, "ego_vehicle")) for c in contours_blue: all_contours.append((c, "other_vehicle")) vehicle_states = [] # 6. Iterate through each detected contour for contour, class_name in all_contours: # Filter out very small contours that might be noise if cv2.contourArea(contour) < 50: continue # --- State Estimation --- # a. Get the minimum area rotated rectangle # This is perfect for finding the orientation of non-upright rectangles. rect = cv2.minAreaRect(contour) (pos_x, pos_y), _, _ = rect # Get the 4 corners of the rotated rectangle for drawing and heading calculation box_points = cv2.boxPoints(rect) box_points = np.intp(box_points) # b. Heading (Robust Calculation) # We find the longer side of the rectangle and calculate its angle. edge1 = np.linalg.norm(box_points[0] - box_points[1]) edge2 = np.linalg.norm(box_points[1] - box_points[2]) # Determine the vector corresponding to the vehicle's length (the longer side) if edge1 > edge2: delta_x = box_points[1][0] - box_points[0][0] delta_y = box_points[1][1] - box_points[0][1] else: delta_x = box_points[2][0] - box_points[1][0] delta_y = box_points[2][1] - box_points[1][1] # Calculate the angle of this vector angle_rad = math.atan2(delta_y, delta_x) heading = math.degrees(angle_rad) # As all vehicles in highway-env move to the right, we ensure the # heading is in the right-hand plane (between -90 and 90 degrees). if heading > 90: heading -= 180 elif heading < -90: heading += 180 # c. Speed # Speed calculation requires tracking across multiple frames. # Since we only have one frame, we'll set it to 0. speed = 0.0 # Placeholder # Store the state vehicle_states.append({ "class": class_name, "bounding_box_points": box_points.tolist(), "position_x": pos_x, "position_y": pos_y, "speed": speed, "heading": heading }) # --- Visualization --- # Draw the rotated bounding box cv2.drawContours(annotated_img, [box_points], 0, (0, 255, 255), 2) # Yellow box # Draw the center point cv2.circle(annotated_img, (int(pos_x), int(pos_y)), 5, (0, 0, 255), -1) # Red dot # Draw the heading vector length = 40 # Length of the heading line angle_rad_viz = np.deg2rad(heading) # Use the corrected heading for visualization end_x = int(pos_x + length * np.cos(angle_rad_viz)) end_y = int(pos_y + length * np.sin(angle_rad_viz)) cv2.line(annotated_img, (int(pos_x), int(pos_y)), (end_x, end_y), (255, 0, 0), 2) # Blue line # Put text label label = f"H: {heading:.1f}" cv2.putText(annotated_img, label, (box_points[1][0], box_points[1][1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2) return annotated_img, vehicle_states # Legacy function for backward compatibility def estimate_vehicle_states_by_color(image_path: str) -> Tuple[Optional[np.ndarray], List[Dict]]: """ Legacy function for backward compatibility. Creates a default model instance and processes the image. """ model = ContourDetectionModel() return model.estimate_vehicle_states_by_color(image_path) def save_states_to_csv(states: List[Dict], csv_file_path: str) -> None: """ Legacy function for backward compatibility. """ model = ContourDetectionModel() model.save_states_to_csv(states, csv_file_path) if __name__ == '__main__': # Example usage when run as script # You can modify these paths as needed input_image_path = '/home/alienware3/Documents/diamond/frames/frame_0.png' output_image_path = 'frame_000001_contour_detected.png' output_csv_path = 'vehicle_states.csv' # Create model instance model = ContourDetectionModel() img = cv2.imread(input_image_path) # Process the image annotated_image, states = model.detect_from_image( img ) cv2.imwrite(output_image_path, annotated_image)