tarinmodel9 / app.py
nagasurendra's picture
Update app.py
04f4d0b verified
raw
history blame
2.59 kB
import cv2
import torch
import gradio as gr
import numpy as np
from ultralytics import YOLO
# Load YOLOv8 model and set device (GPU if available)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = YOLO('./data/best.pt') # Path to your model
model.to(device)
# Define the function that processes the uploaded video
def process_video(video):
# video is now the file path string, not a file object
input_video = cv2.VideoCapture(video) # Directly pass the path to cv2.VideoCapture
# Get frame width, height, and fps from input video
frame_width = int(input_video.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(input_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = input_video.get(cv2.CAP_PROP_FPS)
# Resize to reduce computation (optional)
new_width, new_height = 640, 480 # Resize to 640x480 resolution
frame_width, frame_height = new_width, new_height
while True:
# Read a frame from the video
ret, frame = input_video.read()
if not ret:
break # End of video
# Resize the frame to reduce computational load
frame = cv2.resize(frame, (new_width, new_height))
# Perform inference on the frame
results = model(frame) # Automatically uses GPU if available
# Check if any object was detected
if len(results[0].boxes) > 0: # If there are detected objects
# Annotate the frame with bounding boxes
annotated_frame = results[0].plot() # Plot the frame with bounding boxes
# Convert the annotated frame to RGB format for displaying
annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
# Display the frame with detections
cv2.imshow("Detected Frame", annotated_frame_rgb)
# Wait for a key press (optional: press 'q' to quit early)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# Release resources
input_video.release()
cv2.destroyAllWindows()
return "Video processing complete!"
# Create a Gradio interface for video upload
iface = gr.Interface(fn=process_video,
inputs=gr.Video(label="Upload Video"), # Updated line
outputs=gr.Textbox(label="Processing Status"), # Output text showing processing status
title="YOLOv8 Object Detection - Real-Time Display",
description="Upload a video for object detection using YOLOv8. The frames with detections will be shown in real-time.")
# Launch the interface
iface.launch()