File size: 2,876 Bytes
eb285c2 7433600 eb285c2 7433600 eb285c2 7433600 eb285c2 7433600 eb285c2 7433600 eb285c2 |
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 |
import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
# Load the YOLOv8 model
model = YOLO("./model/best.pt")
def detect_emotion(image):
"""
Perform YOLO8 inference on the uploaded image.
:param image: Input image from the Gradio interface
:return: Annotated image with bounding boxes and emotion labels
"""
# Convert PIL image to OpenCV format
image = np.array(image)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# Perform inference
results = model(image)
# Annotate the image with predictions
annotated_image = results[0].plot()
# Convert OpenCV BGR image back to RGB for display
annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
return annotated_image
def detect_emotion_video(video_path):
"""
Perform YOLO8 inference on an uploaded video.
:param video_path: Path to the video file from Gradio interface
:return: Processed video with bounding boxes and emotion labels
"""
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return "Error: Could not open video file."
# Get video properties
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))
fps = int(cap.get(cv2.CAP_PROP_FPS))
# Define the output video writer
output_video_path = "output_video.mp4"
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break # Stop if video ends
# Perform inference
results = model(frame)
# Annotate the frame with predictions
annotated_frame = results[0].plot()
# Write the processed frame to the output video
out.write(annotated_frame)
cap.release()
out.release()
return output_video_path # Return the processed video
# Create Gradio Tabs
with gr.Blocks() as demo:
gr.Markdown("## YOLOv8 Fruits Detection")
with gr.Tabs():
# Tab 1: Image Inference
with gr.Tab("Image Detection"):
gr.Markdown("### Upload an Image for Fruits Detection")
image_input = gr.Image(type="pil")
image_output = gr.Image(type="numpy")
image_btn = gr.Button("Detect Fruit")
image_btn.click(detect_emotion, inputs=image_input, outputs=image_output)
# Tab 2: Video Inference
with gr.Tab("Video Detection"):
gr.Markdown("### Upload a Video for Fruits Detection")
video_input = gr.Video()
video_output = gr.Video()
video_btn = gr.Button("Detect Fruits in Video")
video_btn.click(detect_emotion_video, inputs=video_input, outputs=video_output)
# Launch the Gradio App
demo.launch(share=True) # Enables public sharing
|