Revert back changes

app.py

@@ -2,78 +2,35 @@ import torch
 import gradio as gr
 from PIL import Image
 import io
-import cv2
-import numpy as np
-import tempfile
 
 # Load the YOLOv5 model
-model = torch.hub.load('ultralytics/yolov5', 'custom', path='fire.pt', trust_repo=True)  # Load custom model
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='fire.pt')  # Load custom model
 
 def detect_objects(image):
-    # Run the YOLOv5 model on an image
+    # Run the YOLOv5 model
     results = model(image)
-    results_image = results.render()[0]  # Render results as a list, take the first element
-    results_image = Image.fromarray(results_image)
-    
-    return results_image
 
-def process_video(video_path):
-    # Read video frames
-    cap = cv2.VideoCapture(video_path)
-    
-    # Get video information
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec
-    
-    # Temporary file to store the output video
-    temp_output = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
-    
-    out = cv2.VideoWriter(temp_output.name, fourcc, fps, (width, height))
+    # Save the results image
+    results_image = results.render()[0]  # Render returns a list, we take the first element
 
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        
-        # Convert frame (BGR to RGB)
-        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        
-        # Run YOLOv5 on the frame
-        results = model(rgb_frame)
-        
-        # Render the results
-        results_frame = results.render()[0]
-        
-        # Convert back to BGR for OpenCV
-        results_frame_bgr = cv2.cvtColor(results_frame, cv2.COLOR_RGB2BGR)
-        
-        # Write frame to output video
-        out.write(results_frame_bgr)
-    
-    cap.release()
-    out.release()
-    
-    return temp_output.name
+    # Convert the numpy array result to an image
+    results_image = Image.fromarray(results_image)
 
-# Gradio interface function
-def detect(input_image=None, input_video=None):
-    if input_video is not None:  # If a video is uploaded
-        return process_video(input_video), None  # Return the video, no image
-    elif input_image is not None:  # If an image is uploaded
-        return None, detect_objects(input_image)  # Return the image, no video
-    else:
-        return None, None
+    # Save to a buffer
+    buf = io.BytesIO()
+    results_image.save(buf, format='JPEG')
+    byte_im = buf.getvalue()
+
+    return results_image
 
 # Gradio interface
 interface = gr.Interface(
-    fn=detect,
-    inputs=[gr.Image(type="pil", label="Upload Image"), gr.Video(label="Upload Video")],  # Two separate inputs
-    outputs=[gr.Video(label="Processed Video"), gr.Image(type="pil", label="Processed Image")],
-    title="YOLOv5 Image & Video Detection",
-    description="Upload an image or video to detect objects using YOLOv5."
+    fn=detect_objects,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="YOLOv5 Image Detection",
+    description="Upload an image to detect objects using YOLOv5."
 )
 
 # Launch the Gradio app
-interface.launch()
+interface.launch()