Ultralytics/YOLOv8
19
1# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license2 3import tempfile4import cv25import gradio as gr6import numpy as np7import PIL.Image as Image8from ultralytics import YOLO9from pathlib import Path10 11 12MODEL_CHOICES = [13 "yolov8n",14 "yolov8s",15 "yolov8m",16 "yolov8n-seg",17 "yolov8s-seg",18 "yolov8m-seg",19 "yolov8n-pose",20 "yolov8s-pose",21 "yolov8m-pose",22 "yolov8n-obb",23 "yolov8s-obb",24 "yolov8m-obb",25 "yolov8n-cls",26 "yolov8s-cls",27 "yolov8m-cls",28]29 30IMAGE_SIZE_CHOICES = [320, 640, 1024]31CUSTOM_CSS = (Path(__file__).parent / "ultralytics.css").read_text()32 33def predict_image(img, conf_threshold, iou_threshold, model_name, show_labels, show_conf, imgsz):34 """Predicts objects in an image using a Ultralytics YOLO model with adjustable confidence and IOU thresholds."""35 model = YOLO(model_name)36 results = model.predict(37 source=img,38 conf=conf_threshold,39 iou=iou_threshold,40 imgsz=imgsz,41 verbose=False,42 )43 44 for r in results:45 im_array = r.plot(labels=show_labels, conf=show_conf)46 im = Image.fromarray(im_array[..., ::-1])47 48 return im49 50 51def predict_video(video_path, conf_threshold, iou_threshold, model_name, show_labels, show_conf, imgsz):52 """Predicts objects in a video using a Ultralytics YOLO model and returns the annotated video."""53 if video_path is None:54 return None55 56 model = YOLO(model_name)57 58 # Open the video59 cap = cv2.VideoCapture(video_path)60 if not cap.isOpened():61 return None62 63 # Get video properties64 fps = int(cap.get(cv2.CAP_PROP_FPS))65 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))66 height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))67 68 # Create temporary output file69 temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)70 output_path = temp_output.name71 temp_output.close()72 73 # Initialize video writer74 fourcc = cv2.VideoWriter_fourcc(*"mp4v")75 out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))76 77 while True:78 ret, frame = cap.read()79 if not ret:80 break81 82 # Run inference on the frame83 results = model.predict(84 source=frame,85 conf=conf_threshold,86 iou=iou_threshold,87 imgsz=imgsz,88 verbose=False,89 )90 91 # Get the annotated frame92 annotated_frame = results[0].plot(labels=show_labels, conf=show_conf)93 out.write(annotated_frame)94 95 cap.release()96 out.release()97 98 return output_path99 100# Cache model for streaming performance101_model_cache = {}102 103def get_model(model_name):104 """Get or create a cached model instance."""105 if model_name not in _model_cache:106 _model_cache[model_name] = YOLO(model_name)107 return _model_cache[model_name]108 109 110def predict_webcam(frame, conf_threshold, iou_threshold, model_name, show_labels, show_conf, imgsz):111 """Predicts objects in a webcam frame using a Ultralytics YOLO model (optimized for streaming)."""112 if frame is None:113 return None114 115 # Use cached model for better streaming performance116 model = get_model(model_name)117 118 if isinstance(frame, np.ndarray):119 # Gradio webcam sends RGB, but Ultralytics YOLO expects BGR for OpenCV operations120 # Convert RGB to BGR for YOLO121 frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)122 123 # Run inference124 results = model.predict(125 source=frame_bgr,126 conf=conf_threshold,127 iou=iou_threshold,128 imgsz=imgsz,129 verbose=False,130 )131 132 # YOLO's plot() returns BGR, convert back to RGB for Gradio display133 annotated_frame = results[0].plot(labels=show_labels, conf=show_conf)134 # Convert BGR to RGB for Gradio135 return cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)136 137 return None138 139 140# Create the Gradio app with tabs141with gr.Blocks(title="Ultralytics YOLOv8 Inference 🚀") as demo:142 gr.Markdown("# Ultralytics YOLOv8 Inference 🚀")143 gr.Markdown("Upload images, videos, or use your webcam for real-time detection, segmentation, pose estimation, OBB, and classification.")144 145 with gr.Tabs():146 # Image Tab147 with gr.TabItem("📷 Image"):148 with gr.Row():149 with gr.Column():150 img_input = gr.Image(type="pil", label="Upload Image")151 img_conf = gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold")152 img_iou = gr.Slider(minimum=0, maximum=1, value=0.7, label="IoU threshold")153 img_model = gr.Radio(choices=MODEL_CHOICES, label="Model Name", value="yolov8n")154 img_labels = gr.Checkbox(value=True, label="Show Labels")155 img_conf_show = gr.Checkbox(value=True, label="Show Confidence")156 img_size = gr.Radio(choices=IMAGE_SIZE_CHOICES, label="Image Size", value=640)157 img_btn = gr.Button("Detect Objects", variant="primary")158 with gr.Column():159 img_output = gr.Image(type="pil", label="Result")160 161 img_btn.click(162 predict_image,163 inputs=[img_input, img_conf, img_iou, img_model, img_labels, img_conf_show, img_size],164 outputs=img_output,165 )166 167 gr.Examples(168 examples=[169 ["https://ultralytics.com/images/bus.jpg", 0.25, 0.7, "yolov8n", True, True, 640],170 ["https://ultralytics.com/images/zidane.jpg", 0.25, 0.7, "yolov8n-seg", True, True, 640],171 ["https://ultralytics.com/images/boats.jpg", 0.25, 0.7, "yolov8n-obb", True, True, 1024],172 ],173 inputs=[img_input, img_conf, img_iou, img_model, img_labels, img_conf_show, img_size],174 )175 176 # Video Tab177 with gr.TabItem("🎬 Video"):178 with gr.Row():179 with gr.Column():180 vid_input = gr.Video(label="Upload Video")181 vid_conf = gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold")182 vid_iou = gr.Slider(minimum=0, maximum=1, value=0.7, label="IoU threshold")183 vid_model = gr.Radio(choices=MODEL_CHOICES, label="Model Name", value="yolov8n")184 vid_labels = gr.Checkbox(value=True, label="Show Labels")185 vid_conf_show = gr.Checkbox(value=True, label="Show Confidence")186 vid_size = gr.Radio(choices=IMAGE_SIZE_CHOICES, label="Image Size", value=640)187 vid_btn = gr.Button("Process Video", variant="primary")188 with gr.Column():189 vid_output = gr.Video(label="Result")190 191 vid_btn.click(192 predict_video,193 inputs=[vid_input, vid_conf, vid_iou, vid_model, vid_labels, vid_conf_show, vid_size],194 outputs=vid_output,195 )196 197 # Webcam Tab - Real-time streaming198 with gr.TabItem("📹 Webcam"):199 gr.Markdown("### Real-time Webcam Detection")200 gr.Markdown("Enable streaming for live detection as you move!")201 with gr.Row():202 with gr.Column():203 webcam_conf = gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold")204 webcam_iou = gr.Slider(minimum=0, maximum=1, value=0.7, label="IoU threshold")205 webcam_model = gr.Radio(choices=MODEL_CHOICES, label="Model Name", value="yolov8n")206 webcam_labels = gr.Checkbox(value=True, label="Show Labels")207 webcam_conf_show = gr.Checkbox(value=True, label="Show Confidence")208 webcam_size = gr.Radio(choices=IMAGE_SIZE_CHOICES, label="Image Size", value=640)209 with gr.Column():210 # Streaming webcam input with real-time output211 webcam_input = gr.Image(212 sources=["webcam"],213 type="numpy",214 label="Webcam (streaming)",215 streaming=True,216 )217 webcam_output = gr.Image(type="numpy", label="Detection Result")218 219 # Stream event for real-time detection220 webcam_input.stream(221 predict_webcam,222 inputs=[223 webcam_input,224 webcam_conf,225 webcam_iou,226 webcam_model,227 webcam_labels,228 webcam_conf_show,229 webcam_size,230 ],231 outputs=webcam_output,232 )233 234demo.launch(css=CUSTOM_CSS, ssr_mode=False)235