Update app.py

app.py

@@ -2,58 +2,24 @@ from ultralytics import YOLO
 import gradio as gr
 from PIL import Image
 import numpy as np
-import cv2
-import tempfile
-import os
 
 # Load the trained YOLOv8 model for human detection
-model = YOLO("best.pt")  # Replace with your trained model path
+model = YOLO("best.pt")  # Replace with the correct path to your human detection model
 
-# Define function to handle both images and videos
-def detect_humans(input_file):
-    # If input is a PIL image (from Gradio image upload)
-    if isinstance(input_file, Image.Image):
-        results = model(input_file)
-        result_img = results[0].plot()
-        return Image.fromarray(result_img.astype(np.uint8))
-    
-    # If input is a video file path
-    elif isinstance(input_file, str) and os.path.isfile(input_file):
-        cap = cv2.VideoCapture(input_file)
-        
-        # Video writer setup
-        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-        temp_output = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
-        out = cv2.VideoWriter(temp_output.name, fourcc, int(cap.get(cv2.CAP_PROP_FPS)),
-                              (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
-                               int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))
-        
-        while True:
-            ret, frame = cap.read()
-            if not ret:
-                break
-            results = model(frame)
-            annotated_frame = results[0].plot()
-            out.write(annotated_frame)
-        
-        cap.release()
-        out.release()
-        
-        return temp_output.name  # Return path to processed video file
-    
-    else:
-        return None
+# Define prediction function
+def detect_humans(image):
+    results = model(image)  # Perform inference
+    result_img = results[0].plot()  # Draw bounding boxes
+    return Image.fromarray(result_img.astype(np.uint8))  # Return image with boxes
 
-# Build Gradio interface (accepts both images and videos)
+# Build Gradio interface
 interface = gr.Interface(
     fn=detect_humans,
-    inputs=gr.File(type="file", label="Upload Image or Video"),
-    outputs=[
-        gr.Image(type="pil", label="Detected Humans (Image)"),
-        gr.Video(label="Detected Humans (Video)")
-    ],
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Image(type="pil", label="Detected Humans"),
     title="Human Detection (YOLOv8)",
-    description="Upload an image or video. The model will detect humans using your trained YOLOv8 model."
+    description="Upload an image. The model will detect humans using your trained YOLOv8 model."
 )
 
+# Launch the interface
 interface.launch(debug=True)