plot_csv_rectangles.py

import pandas as pd
import cv2
import numpy as np
import random

def plot_mot_rectangles_opencv(csv_file, image_name, output_image_path="mot_rectangles.png"):
    """
    Plots rectangles from a MOTChallenge CSV file in a single image,
    with different colors for each identity, and saves it using OpenCV.

    Args:
        csv_file (str): Path to the MOTChallenge CSV file.
        output_image_path (str): Path to save the output image (e.g., "mot_rectangles.png").
    """

    try:
        df = pd.read_csv(csv_file,sep=',')
    except FileNotFoundError:
        print(f"Error: File not found: {csv_file}")
        return

    df.columns = ['frame','id','bb_left','bb_top','bb_width','bb_height','conf','x','y','z']
    unique_ids = df['id'].unique()
    id_colors = {id: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for id in unique_ids}

    max_x = int(df['bb_left'].max() + df['bb_width'].max())
    max_y = int(df['bb_top'].max() + df['bb_height'].max())

    # image = np.zeros((max_y, max_x, 3), dtype=np.uint8)  # Create a black image
    image = cv2.imread(image_name)
    
    for _, row in df.iterrows():
        x, y, width, height, identity = int(row['bb_left']), int(row['bb_top']), int(row['bb_width']), int(row['bb_height']), row['id']
        color = id_colors[identity]
        cv2.rectangle(image, (x, y), (x + width, y + height), color, 2)

    cv2.imwrite(output_image_path, image)
    print(f"Image saved to: {output_image_path}")



# Example Usage:
plot_mot_rectangles_opencv("annotation_example_straight.csv", "data/image.jpg", "output.png")