add support for image inference

app.py

@@ -1,12 +1,11 @@
 import gradio as gr
 import PIL.Image as Image
-import tempfile
-import cv2
+
 from ultralytics import ASSETS, YOLO
 
-# Load YOLOv8 model
 model = YOLO("yolov8n.pt")
 
+
 def predict_image(img, conf_threshold, iou_threshold):
     """Predicts objects in an image using a YOLOv8 model with adjustable confidence and IOU thresholds."""
     results = model.predict(
@@ -24,74 +23,19 @@ def predict_image(img, conf_threshold, iou_threshold):
 
     return im
 
-def predict_video(video_path, conf_threshold, iou_threshold):
-    """Predicts objects in a video using a YOLOv8 model with adjustable confidence and IOU thresholds."""
-    # Create a temporary file to save the processed video
-    temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
-    temp_output.close()
-
-    # Load video
-    cap = cv2.VideoCapture(video_path)
-
-    # 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))
-
-    # Set up VideoWriter to save output video
-    out = cv2.VideoWriter(temp_output.name, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
-
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-
-        # Perform inference on each frame
-        results = model.predict(
-            source=frame,
-            conf=conf_threshold,
-            iou=iou_threshold,
-            show_labels=True,
-            show_conf=True,
-            imgsz=640,
-        )
-
-        # Draw the results on the frame
-        for r in results:
-            frame = r.plot()
-
-        # Write the frame to the output video
-        out.write(frame)
-
-    # Release resources
-    cap.release()
-    out.release()
-
-    return temp_output.name
-
-def process_input(input_file, conf_threshold, iou_threshold, mode):
-    """Handles both image and video inference based on the selected mode."""
-    if mode == "Image":
-        return predict_image(input_file, conf_threshold, iou_threshold)
-    elif mode == "Video":
-        return predict_video(input_file.name, conf_threshold, iou_threshold)
-
-# Create Gradio interface
 iface = gr.Interface(
-    fn=process_input,
+    fn=predict_image,
     inputs=[
-        gr.File(label="Upload Image or Video File"),  # Use a generic File input for both image and video
+        gr.Image(type="pil", label="Upload Image"),
         gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
         gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
-        gr.Radio(choices=["Image", "Video"], label="Select Mode", value="Image"),
     ],
-    outputs=gr.Image(type="pil", label="Result") if gr.Radio else gr.Video(type="file", label="Result"),
+    outputs=gr.Image(type="pil", label="Result"),
     title="Ultralytics Gradio Application 🚀",
-    description="Upload images or videos for inference. The Ultralytics YOLOv8n model is used by default.",
+    description="Upload images for inference. The Ultralytics YOLOv8n model is used by default.",
     examples=[
-        [ASSETS / "bus.jpg", 0.25, 0.45, "Image"],
-        [ASSETS / "zidane.jpg", 0.25, 0.45, "Image"],
+        [ASSETS / "bus.jpg", 0.25, 0.45],
+        [ASSETS / "zidane.jpg", 0.25, 0.45],
     ],
 )
-
 iface.launch(share=True)