app.py

import cv2
import torch
import gradio as gr
import numpy as np
import os
import matplotlib.pyplot as plt
from ultralytics import YOLO, __version__ as ultralytics_version
import uuid

# Debug: Check environment
print(f"Torch version: {torch.__version__}")
print(f"Gradio version: {gr.__version__}")
print(f"Ultralytics version: {ultralytics_version}")
print(f"CUDA available: {torch.cuda.is_available()}")

# Load YOLOv8 model
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
model = YOLO('./data/best.pt').to(device)

def process_video(video, output_folder="detected_frames", plot_graphs=False):
    if video is None:
        yield "Error: No video uploaded", []
        return
    
    # Create output folder if it doesn't exist
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    cap = cv2.VideoCapture(video)
    if not cap.isOpened():
        yield "Error: Could not open video file", []
        return
    
    frame_width, frame_height = 320, 240  # Smaller resolution
    frame_count = 0
    frame_skip = 5  # Process every 5th frame
    max_frames = 100  # Limit for testing
    confidence_scores = []  # Store confidence scores for plotting
    detected_frame_paths = []  # Store paths of frames with detections
    
    while True:
        ret, frame = cap.read()
        if not ret or frame_count > max_frames:
            break
            
        frame_count += 1
        if frame_count % frame_skip != 0:
            continue
            
        frame = cv2.resize(frame, (frame_width, frame_height))
        print(f"Processing frame {frame_count}")
        
        # Run YOLOv8 inference
        results = model(frame)
        
        # Save and yield frame if objects are detected
        if results[0].boxes is not None and len(results[0].boxes) > 0:
            annotated_frame = results[0].plot()
            frame_filename = os.path.join(output_folder, f"frame_{frame_count:04d}.jpg")
            cv2.imwrite(frame_filename, annotated_frame)
            detected_frame_paths.append(frame_filename)
            
            # Collect confidence scores for plotting
            confs = results[0].boxes.conf.cpu().numpy()
            confidence_scores.extend(confs)
            
            # Yield current status and gallery
            yield f"Processed frame {frame_count} with detections", detected_frame_paths[:]
    
    cap.release()
    
    # Generate confidence score plot if requested
    if plot_graphs and confidence_scores:
        plt.figure(figsize=(10, 5))
        plt.hist(confidence_scores, bins=20, color='blue', alpha=0.7)
        plt.title('Distribution of Confidence Scores')
        plt.xlabel('Confidence Score')
        plt.ylabel('Frequency')
        graph_path = os.path.join(output_folder, "confidence_histogram.png")
        plt.savefig(graph_path)
        plt.close()
        detected_frame_paths.append(graph_path)
    
    # Final yield with all results
    status = f"Saved {len(detected_frame_paths)} frames with detections in {output_folder}. {f'Graph saved as {graph_path}' if plot_graphs and confidence_scores else ''}"
    yield status, detected_frame_paths

# Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# YOLOv8 Object Detection - Real-time Frame Output")
    gr.Markdown("Upload a short video to view frames with detections immediately in a gallery. Optionally generate a confidence score graph.")
    
    with gr.Row():
        video_input = gr.Video(label="Upload Video")
        output_folder = gr.Textbox(label="Output Folder", value="detected_frames")
        plot_graphs = gr.Checkbox(label="Generate Confidence Score Graph", value=False)
    
    submit_button = gr.Button("Process Video")
    
    status_output = gr.Text(label="Status")
    gallery_output = gr.Gallery(label="Detected Frames and Graph", preview=True, columns=3)
    
    submit_button.click(
        fn=process_video,
        inputs=[video_input, output_folder, plot_graphs],
        outputs=[status_output, gallery_output],
        concurrency_limit=1
    )

if __name__ == "__main__":
    iface.launch()