Spaces:
Running
Running
File size: 3,131 Bytes
8d18811 783e9d0 8d18811 783e9d0 8d18811 637e0b8 783e9d0 637e0b8 783e9d0 637e0b8 8d18811 637e0b8 8d18811 637e0b8 8d18811 637e0b8 8d18811 637e0b8 e42a8c2 783e9d0 8d18811 637e0b8 8d18811 783e9d0 c35ff0b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
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(
source=img,
conf=conf_threshold,
iou=iou_threshold,
show_labels=True,
show_conf=True,
imgsz=640,
)
for r in results:
im_array = r.plot()
im = Image.fromarray(im_array[..., ::-1])
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,
inputs=[
gr.File(label="Upload Image or Video File"), # Use a generic File input for both image and video
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"),
title="Ultralytics Gradio Application 🚀",
description="Upload images or videos 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"],
],
)
iface.launch(share=True)
|