Update app.py

app.py

@@ -1,88 +1,26 @@
 from ultralytics import YOLO
 from PIL import Image
 import gradio as gr
-import cv2
-import tempfile
-import os
 
 # Load YOLOv8 model
 model = YOLO("best.pt")  # Ensure best.pt is in the same directory
 
-# Preprocess and run inference for images
-def predict_image(image):
+# Preprocess and run inference
+def predict(image):
+    # Perform prediction
     results = model.predict(source=image, conf=0.5)
+    
+    # Annotate the image with bounding boxes
     annotated_image = results[0].plot()
+    
+    # Convert to PIL Image
     return Image.fromarray(annotated_image)
 
-# Preprocess and run inference for videos
-def predict_video(video):
-    try:
-        # Save video to a temporary file
-        temp_video_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
-        with open(temp_video_path, "wb") as f:
-            f.write(video.read())
-
-        # Open the video file
-        cap = cv2.VideoCapture(temp_video_path)
-        if not cap.isOpened():
-            return "Error: Unable to open video file."
-
-        # Get video properties
-        fps = int(cap.get(cv2.CAP_PROP_FPS))
-        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
-        # Define codec and create a video writer
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Use 'mp4v' for compatibility
-        output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
-        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-
-            # Perform predictions on the frame
-            results = model.predict(source=frame, conf=0.5)
-            annotated_frame = results[0].plot()
-
-            # Write the annotated frame
-            out.write(annotated_frame)
-
-        # Release resources
-        cap.release()
-        out.release()
-
-        # Ensure the output file exists and is playable
-        if os.path.exists(output_path):
-            return output_path
-        else:
-            return "Error: Annotated video could not be created."
-
-    except Exception as e:
-        return f"An error occurred while processing the video: {str(e)}"
-
-# Gradio interfaces
-image_interface = gr.Interface(
-    fn=predict_image,
-    inputs=gr.Image(type="pil", label="Upload an Image"),
+# Gradio interface
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
     outputs="image",
-    title="Image Detection",
-    description="Upload an image for object detection."
-)
-
-video_interface = gr.Interface(
-    fn=predict_video,
-    inputs=gr.Video(label="Upload a Video"),
-    outputs="video",
-    title="Video Detection",
-    description="Upload a video for object detection."
-)
-
-# Combine interfaces
-app = gr.TabbedInterface(
-    [image_interface, video_interface],
-    ["Image Detection", "Video Detection"]
-)
-
-app.launch()
+    title="Hippo or Rhino Detection",
+    description="Upload an image for object detection with YOLOv8."
+).launch()