app.py

import gradio as gr
from ultralytics import YOLO
from PIL import Image
import cv2
import numpy as np
import tempfile
import os

# Load the YOLOv8 model
model = YOLO('yolov8n.pt')

def process_image(image):
    results = model(image)
    # Get detection information
    boxes = results[0].boxes
    detection_info = []
    for box in boxes:
        class_id = int(box.cls[0])
        class_name = results[0].names[class_id]
        confidence = float(box.conf[0])
        detection_info.append(f"{class_name}: {confidence:.2%}")
    
    return Image.fromarray(results[0].plot()), "\n".join(detection_info)

def process_video(video_path):
    with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_file:
        output_path = temp_file.name

    cap = cv2.VideoCapture(video_path)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    detection_summary = []
    frame_count = 0
    
    try:
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
                
            frame_count += 1
            results = model(frame)
            
            # Collect detection information for this frame
            if frame_count % int(fps) == 0:  # Sample every second
                for box in results[0].boxes:
                    class_id = int(box.cls[0])
                    class_name = results[0].names[class_id]
                    detection_summary.append(class_name)
            
            annotated_frame = results[0].plot()
            out.write(annotated_frame)
            
    finally:
        cap.release()
        out.release()
    
    # Create summary of detected objects
    if detection_summary:
        from collections import Counter
        counts = Counter(detection_summary)
        summary = "\n".join([f"{obj}: {count} occurrences" for obj, count in counts.most_common()])
    else:
        summary = "No objects detected"
    
    return output_path, summary

def detect_objects(media):
    if media is None:
        return None, None, None, "Please upload an image or video to begin detection.", gr.update(visible=True), gr.update(visible=False)
    
    try:
        if isinstance(media, str) and media.lower().endswith(('.mp4', '.avi', '.mov')):
            output_video, detection_summary = process_video(media)
            return (None, output_video, detection_summary,
                    "✅ Video processing complete! Check the detection summary below.", 
                    gr.update(visible=False), gr.update(visible=True))
        else:
            if isinstance(media, str):
                image = cv2.imread(media)
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            else:
                image = media
            processed_image, detection_info = process_image(image)
            return (processed_image, None, detection_info,
                    "✅ Image processing complete! Check the detections below.", 
                    gr.update(visible=True), gr.update(visible=False))
    except Exception as e:
        return None, None, None, f"❌ Error: {str(e)}", gr.update(visible=False), gr.update(visible=False)

# Custom CSS for styling
custom_css = """
#app-container {
    max-width: 1200px;
    margin: 0 auto;
    padding: 20px;
}

#logo-img {
    display: block;
    margin: 0 auto;
    max-height: 100px;
    margin-bottom: 20px;
}

.upload-box {
    border: 2px dashed #ccc;
    padding: 20px;
    text-align: center;
    border-radius: 8px;
    background-color: #f8f9fa;
    margin: 20px 0;
}

.results-container {
    background-color: #ffffff;
    border-radius: 8px;
    padding: 15px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    margin-top: 20px;
}

.detection-info {
    background-color: #f8f9fa;
    padding: 15px;
    border-radius: 8px;
    margin-top: 10px;
    font-family: monospace;
}
"""

# Create Gradio interface
with gr.Blocks(css=custom_css) as demo:
    with gr.Column(elem_id="app-container"):
        # Logo and Header
        gr.HTML(
            """
            <div style="text-align: center; margin-bottom: 1rem">
                <img src="logo-h.png" id="logo-img" alt="Logo">
            </div>
            """
        )
        
        with gr.Column():
            gr.Markdown("# 🔍 Object Detection")
            
            # Upload Section
            with gr.Column(elem_classes="upload-box"):
                gr.Markdown("### 📤 Upload your file")
                input_media = gr.File(
                    label="Drag and drop or click to upload (Images: jpg, jpeg, png | Videos: mp4, avi, mov)",
                    file_types=["image", "video"]
                )
            
            # Status Message
            status_text = gr.Textbox(
                label="Status",
                value="Waiting for upload...",
                interactive=False
            )
            
            # Detection Information
            detection_info = gr.Textbox(
                label="Detection Results",
                elem_classes="detection-info",
                interactive=False
            )
        
        # Results Section
        with gr.Column(elem_classes="results-container"):
            with gr.Row():
                with gr.Column(visible=False) as image_column:
                    output_image = gr.Image(label="Detected Objects")
                with gr.Column(visible=False) as video_column:
                    output_video = gr.Video(label="Processed Video")
    
    # Handle file upload
    input_media.upload(
        fn=detect_objects,
        inputs=[input_media],
        outputs=[output_image, output_video, detection_info, status_text, 
                image_column, video_column]
    )

if __name__ == "__main__":
    demo.launch(share=True)