src/models/contours.py

# Vehicle Detection and State Estimation using Color-Based Contour Detection
#
# This script detects vehicles in a Bird's-Eye View (BEV) image by
# isolating their specific colors (green and blue) and then analyzing
# the shapes (contours) of the colored areas. The detected states are
# then exported to a CSV file.
#
# 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

import cv2
import numpy as np
import math
import csv # Import the csv module

def estimate_vehicle_states_by_color(image_path):
    """
    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.
    """
    # 1. 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, []

    # 2. Convert the image to HSV color space
    # HSV (Hue, Saturation, Value) is often easier for color segmentation
    # than the default BGR format.
    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

def save_states_to_csv(states, csv_file_path):
    """
    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}")


if __name__ == '__main__':
    # Define input and output file paths
    input_image_path = '/home/zhexiao/Documents/diamond/images/model_output.png'#'/home/zhexiao/Documents/highway_dataset/record_1/images/frame_000165.png'
    output_image_path = 'frame_000001_contour_detected.png'
    output_csv_path = 'vehicle_states.csv'

    # Process the image to get states and the annotated image
    annotated_image, states = estimate_vehicle_states_by_color(input_image_path)

    if annotated_image is not None and states:
        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
        cv2.imwrite(output_image_path, annotated_image)
        print(f"\nAnnotated image saved to: {output_image_path}")
        
        # Save the states to a CSV file
        save_states_to_csv(states, output_csv_path)

    elif not states:
        print("No vehicles were detected in the image.")
    
    # To display the image in a window (if you are running this on a desktop)
    # if annotated_image is not None:
    #     cv2.imshow('Vehicle Detection', annotated_image)
    #     cv2.waitKey(0)
    #     cv2.destroyAllWindows()