Update app.py

app.py

@@ -1,124 +1,100 @@
-import streamlit as st
+import gradio as gr
+import torch
+from ultralytics import YOLO
 import cv2
 import numpy as np
-from ultralytics import YOLO
-from huggingface_hub import hf_hub_download
 from PIL import Image
 import os
-import tempfile
-import supervision as sv
-
-# Title and description
-st.title("DRS Review System - Ball Detection")
-st.write("Upload an image or video to detect balls using a YOLOv5 model for Decision Review System (DRS).")
-
-# Model loading
-@st.cache_resource
-def load_model():
-    # Replace 'your-username/your-repo' with your Hugging Face repository and model file
-    model_path = hf_hub_download(repo_id="your-username/your-repo", filename="best.pt")
-    model = YOLO(model_path)
-    return model
-
-model = load_model()
 
-# Confidence threshold slider
-confidence_threshold = st.slider("Confidence Threshold", 0.0, 1.0, 0.7, 0.05)
+# Load the YOLOv5 model
+model = YOLO("best.pt")
 
-# File uploader for image or video
-uploaded_file = st.file_uploader("Upload an image or video", type=["jpg", "jpeg", "png", "mp4"])
-
-if uploaded_file is not None:
-    # Create a temporary file to save the uploaded content
-    tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.' + uploaded_file.name.split('.')[-1])
-    tfile.write(uploaded_file.read())
-    tfile.close()
-    file_path = tfile.name
-
-    # Check if the uploaded file is an image
-    if uploaded_file.type in ["image/jpeg", "image/png"]:
-        st.subheader("Image Detection Results")
-        image = cv2.imread(file_path)
-        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-        
-        # Run inference
-        results = model(image, conf=confidence_threshold)
-        detections = sv.Detections.from_ultralytics(results[0])
+def detect_ball(input_media, conf_threshold=0.5, iou_threshold=0.5):
+    """
+    Perform ball detection on image or video input.
+    
+    Args:
+        input_media: Uploaded image or video file
+        conf_threshold: Confidence threshold for detection
+        iou_threshold: IoU threshold for non-max suppression
+    
+    Returns:
+        Annotated image or video path
+    """
+    # Check if input is image or video based on file extension
+    file_extension = os.path.splitext(input_media)[1].lower()
+    
+    if file_extension in ['.jpg', '.jpeg', '.png']:
+        # Process image
+        img = cv2.imread(input_media)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         
-        # Annotate image
-        box_annotator = sv.BoxAnnotator()
-        annotated_image = box_annotator.annotate(scene=image_rgb, detections=detections)
+        # Perform detection
+        results = model.predict(img, conf=conf_threshold, iou=iou_threshold)
         
-        # Display result
-        st.image(annotated_image, caption="Detected Balls", use_column_width=True)
-        
-        # Display detection details
-        for score, label, box in zip(detections.confidence, detections.class_id, detections.xyxy):
-            st.write(f"Detected ball with confidence {score:.2f} at coordinates {box.tolist()}")
-
-    # Check if the uploaded file is a video
-    elif uploaded_file.type == "video/mp4":
-        st.subheader("Video Detection Results")
-        output_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
+        # Draw bounding boxes
+        for box in results[0].boxes:
+            x1, y1, x2, y2 = map(int, box.xyxy[0])
+            conf = box.conf[0]
+            label = f"Ball: {conf:.2f}"
+            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
         
+        # Convert to PIL Image for Gradio output
+        output_img = Image.fromarray(img)
+        return output_img
+    
+    elif file_extension in ['.mp4', '.avi', '.mov']:
         # Process video
-        cap = cv2.VideoCapture(file_path)
-        if not cap.isOpened():
-            st.error("Error: Could not open video file.")
-        else:
-            # Get video properties
-            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))
-            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-            
-            # Progress bar
-            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-            progress = st.progress(0)
-            frame_count = 0
-            
-            # Process frames
-            while cap.isOpened():
-                ret, frame = cap.read()
-                if not ret:
-                    break
-                
-                # Run inference on frame
-                results = model(frame, conf=confidence_threshold)
-                detections = sv.Detections.from_ultralytics(results[0])
-                
-                # Annotate frame
-                box_annotator = sv.BoxAnnotator()
-                annotated_frame = box_annotator.annotate(scene=cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), detections=detections)
-                annotated_frame_bgr = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
-                
-                # Write to output video
-                out.write(annotated_frame_bgr)
+        cap = cv2.VideoCapture(input_media)
+        output_path = "output_video.mp4"
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(output_path, fourcc, 30.0, 
+                             (int(cap.get(3)), int(cap.get(4))))
+        
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
                 
-                # Update progress
-                frame_count += 1
-                progress.progress(frame_count / total_frames)
-            
-            cap.release()
-            out.release()
+            # Perform detection
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            results = model.predict(frame_rgb, conf=conf_threshold, iou=iou_threshold)
             
-            # Display video
-            st.video(output_path)
+            # Draw bounding boxes
+            for box in results[0].boxes:
+                x1, y1, x2, y2 = map(int, box.xyxy[0])
+                conf = box.conf[0]
+                label = f"Ball: {conf:.2f}"
+                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
             
-            # Provide download link for processed video
-            with open(output_path, "rb") as file:
-                st.download_button(
-                    label="Download Processed Video",
-                    data=file,
-                    file_name="processed_drs_video.mp4",
-                    mime="video/mp4"
-                )
+            out.write(frame)
         
-        # Clean up temporary files
-        os.remove(file_path)
-        if os.path.exists(output_path):
-            os.remove(output_path)
+        cap.release()
+        out.release()
+        return output_path
+    
+    else:
+        return "Unsupported file format. Please upload an image (.jpg, .png) or video (.mp4, .avi, .mov)."
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Decision Review System (DRS) for Ball Detection")
+    gr.Markdown("Upload an image or video to detect the ball using a trained YOLOv5 model. Adjust confidence and IoU thresholds for detection.")
+    
+ --
+
+    input_media = gr.File(label="Upload Image or Video")
+    conf_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="Confidence Threshold")
+    iou_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="IoU Threshold")
+    output = gr.Image(label="Output (Image or Video)")
+    submit_button = gr.Button("Detect Ball")
+    
+    submit_button.click(
+        fn=detect_ball,
+        inputs=[input_media, conf_slider, iou_slider],
+        outputs=output
+    )
 
-else:
-    st.info("Please upload an image or video to start the DRS review.")
+demo.launch()