main.py

import cv2
import pandas as pd
import pickle
import requests
import matplotlib.pyplot as plt
import argparse
import os


def parse_args():
    parser = argparse.ArgumentParser(description="Process video and extract insights")
    parser.add_argument("--dataset_id", help="Dataset ID (required)")
    parser.add_argument("--version_id", default="1", help="Version ID (default: 1)")
    parser.add_argument("--api_key", help="API key (required)")
    parser.add_argument("--video_path", help="Path to the video (required)")
    parser.add_argument("--interval_minutes", type=int, default=1, help="Interval in seconds (default: 60)")
    return parser.parse_args()


def extract_frames(video_path, interval_minutes):
    cap = cv2.VideoCapture(video_path)
    frames = []
    timestamps = []
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    frame_count = 0

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        if frame_count % (fps * interval_minutes) == 0:
            frames.append(frame)
            timestamps.append(frame_count / fps)
        frame_count += 1
    cap.release()
    return frames, timestamps


def fetch_predictions(base_url, frames, timestamps, dataset_id, version_id, api_key, confidence=0.5):
    headers = {"Content-Type": "application/x-www-form-urlencoded"}
    df_rows = []
    for idx, frame in enumerate(frames):
        numpy_data = pickle.dumps(frame)
        res = requests.post(
            f"{base_url}/{dataset_id}/{version_id}",
            data=numpy_data,
            headers=headers,
            params={"api_key": api_key, "confidence": confidence, "image_type": "numpy"}
        )
        predictions = res.json()

        for pred in predictions['predictions']:
            time_interval = f"{int(timestamps[idx] // 60)}:{int(timestamps[idx] % 60):02}"
            row = {
                "timestamp": time_interval,
                "time": predictions['time'],
                "x": pred["x"],
                "y": pred["y"],
                "width": pred["width"],
                "height": pred["height"],
                "pred_confidence": pred["confidence"],
                "class": pred["class"]
            }
            df_rows.append(row)

    df = pd.DataFrame(df_rows)
    df['seconds'] = df['timestamp'].str.split(':').apply(lambda x: int(x[0])*60 + int(x[1]))
    df = df.sort_values(by="seconds")
    return df

def plot_and_save(data, title, filename, ylabel, stacked=False, legend_title=None, legend_loc=None, legend_bbox=None):
    plt.style.use('dark_background')
    data.plot(kind='bar', stacked=stacked, figsize=(15,7))
    plt.title(title)
    plt.ylabel(ylabel)
    plt.xlabel('Timestamp (in minutes:seconds)')
    
    if legend_title:
        plt.legend(title=legend_title, loc=legend_loc, bbox_to_anchor=legend_bbox)
    
    plt.tight_layout()
    plt.savefig(filename)

def main():
    args = parse_args()
    base_url = "http://localhost:9001"
    video_path = args.video_path
    dataset_id = args.dataset_id
    version_id = args.version_id
    api_key = args.api_key
    interval_minutes = args.interval_minutes * 60


    frames, timestamps = extract_frames(video_path, interval_minutes)
    df = fetch_predictions(base_url, frames, timestamps, dataset_id, version_id, api_key)

    if not os.path.exists("results"):
        os.makedirs("results")

    #saving predictions response to csv
    df.to_csv("results/predictions.csv", index=False)

    # Transform timestamps to minutes and group
    df['minutes'] = df['timestamp'].str.split(':').apply(lambda x: int(x[0]) * 60 + int(x[1]))
    object_counts_per_interval = df.groupby('minutes').size().sort_index()
    object_counts_per_interval.index = object_counts_per_interval.index.map(lambda x: f"{x // 60}:{x % 60:02}")
    object_counts_per_interval.to_csv("results/object_counts_per_interval.csv")

    # Quick insights
    print(f"Total unique objects detected: {df['class'].nunique()}")
    print(f"Most frequently detected object: {df['class'].value_counts().idxmax()}")
    print(f"Time interval with the most objects detected: {object_counts_per_interval.idxmax()}")
    print(f"Time interval with the least objects detected: {object_counts_per_interval.idxmin()}")

    plot_and_save(object_counts_per_interval, 'Number of Objects Detected Over Time', "results/objects_over_time_d.png", 'Number of Objects')

    # Group by timestamp and class, then sort by minutes
    objects_by_class_per_interval = df.groupby(['minutes', 'class']).size().unstack(fill_value=0).sort_index()
    objects_by_class_per_interval.index = objects_by_class_per_interval.index.map(lambda x: f"{x // 60}:{x % 60:02}")
    objects_by_class_per_interval.to_csv("results/object_counts_by_class_per_interval.csv")

    plot_and_save(objects_by_class_per_interval, 'Number of Objects Detected Over Time by Class', "results/objects_by_class_over_time.png", 'Number of Objects', True, "Object Class", "center left", (1, 0.5))


if __name__ == "__main__":
    main()