Spaces:
Sleeping
Sleeping
File size: 1,998 Bytes
1ae06b1 |
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 |
import gradio as gr
from ultralytics import YOLO
from PIL import Image
import cv2
import tempfile
# Load model
model = YOLO("best.pt")
# ---------------- IMAGE DETECTION ----------------
def detect(image):
results = model.predict(image, conf=0.4, iou=0.5)
return results[0].plot()
# ---------------- VIDEO DETECTION ----------------
def detect_video(video):
input_path = video
output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
cap = cv2.VideoCapture(input_path)
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")
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
while True:
ret, frame = cap.read()
if not ret:
break
results = model.predict(frame, conf=0.4, iou=0.5)
r = results[0]
frame = r.plot()
out.write(frame)
cap.release()
out.release()
return output_path
# ---------------- UI ----------------
image_ui = gr.Interface(
fn=detect,
inputs=gr.Image(type="pil", label="Upload Image"),
outputs=gr.Image(label="Detection Result"),
title="Spectacles and Teapots Object Detection - 4124139E",
description="Upload an image or choose any example at the bottom for object detection.",
examples=[
["image_test1.jpg"],
["image_test2.jpg"],
],
cache_examples=False,
)
video_ui = gr.Interface(
fn=detect_video,
inputs=gr.Video(label="Upload Video (.mp4)"),
outputs=gr.Video(label="Processed Video"),
title="Spectacles and Teapots Object Detection - 4124139E",
description="Upload a short video or choose any example at the bottom for object detection.",
examples=[
["harry_potter_short.mp4"]
]
)
demo = gr.TabbedInterface(
[image_ui, video_ui],
["Images", "Video"]
)
demo.launch(server_name="0.0.0.0", server_port=7860) |