Prasanna1622's picture
Update app.py
b83a2f7 verified
raw
history blame
1.75 kB
import gradio as gr
from ultralytics import YOLO
import tempfile
import cv2
# Load YOLOv8 model once
model = "best.pt"
# Inference on image
def predict_image(image):
results = model.predict(image)
return results[0].plot()
# Inference on video
def predict_video(video_path):
# OpenCV video capture
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) // 2 # Resize for performance
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) // 2
# Output video
temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(temp_output.name, fourcc, fps, (width, height))
while True:
ret, frame = cap.read()
if not ret:
break
frame = cv2.resize(frame, (width, height)) # Resize to speed up inference
results = model.predict(frame, imgsz=480, conf=0.5, verbose=False) # Lower imgsz = faster
annotated = results[0].plot()
out.write(annotated)
cap.release()
out.release()
return temp_output.name
# Gradio Interface
with gr.Blocks() as demo:
gr.Markdown("# πŸš€ Optimized YOLOv8 Detection\nFast & Accurate on Images and Videos")
with gr.Tab("Image"):
img_input = gr.Image(type="pil")
img_output = gr.Image(label="Detected")
img_btn = gr.Button("Run Detection")
img_btn.click(predict_image, inputs=img_input, outputs=img_output)
with gr.Tab("Video"):
vid_input = gr.Video()
vid_output = gr.Video()
vid_btn = gr.Button("Run Detection on Video")
vid_btn.click(predict_video, inputs=vid_input, outputs=vid_output)
demo.launch()