Create app.py

app.py

@@ -0,0 +1,206 @@
+import gradio as gr
+from ultralytics import YOLO
+from PIL import Image
+import cv2
+import numpy as np
+import tempfile
+import os
+from pathlib import Path
+
+# Load the YOLOv8 model
+model = YOLO('yolov8n.pt')
+
+def process_image(image):
+    """
+    Process a single image for object detection
+    """
+    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):
+    """
+    Process video for object detection
+    """
+    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):
+    """
+    Unified function to handle both image and video inputs
+    """
+    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 {
+    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;
+}
+.center {
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    margin-bottom: 1rem;
+}
+"""
+
+# Create Gradio interface
+with gr.Blocks(css=custom_css) as demo:
+    with gr.Column(elem_id="app-container"):
+        # Logo and Header
+        with gr.Column(elem_classes="center"):
+            gr.Image("logo-h.png", 
+                    show_label=False, 
+                    container=False,
+                    elem_id="logo-img",
+                    height=100)
+        
+        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)