213 lines
10 KiB
Python
213 lines
10 KiB
Python
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
|
|
|
import io
|
|
from typing import Any, List
|
|
|
|
import cv2
|
|
import torch
|
|
|
|
from ultralytics import YOLO
|
|
from ultralytics.utils import LOGGER
|
|
from ultralytics.utils.checks import check_requirements
|
|
from ultralytics.utils.downloads import GITHUB_ASSETS_STEMS
|
|
|
|
torch.classes.__path__ = [] # Torch module __path__._path issue: https://github.com/datalab-to/marker/issues/442
|
|
|
|
|
|
class Inference:
|
|
"""
|
|
A class to perform object detection, image classification, image segmentation and pose estimation inference.
|
|
|
|
This class provides functionalities for loading models, configuring settings, uploading video files, and performing
|
|
real-time inference using Streamlit and Ultralytics YOLO models.
|
|
|
|
Attributes:
|
|
st (module): Streamlit module for UI creation.
|
|
temp_dict (dict): Temporary dictionary to store the model path and other configuration.
|
|
model_path (str): Path to the loaded model.
|
|
model (YOLO): The YOLO model instance.
|
|
source (str): Selected video source (webcam or video file).
|
|
enable_trk (bool): Enable tracking option.
|
|
conf (float): Confidence threshold for detection.
|
|
iou (float): IoU threshold for non-maximum suppression.
|
|
org_frame (Any): Container for the original frame to be displayed.
|
|
ann_frame (Any): Container for the annotated frame to be displayed.
|
|
vid_file_name (str | int): Name of the uploaded video file or webcam index.
|
|
selected_ind (List[int]): List of selected class indices for detection.
|
|
|
|
Methods:
|
|
web_ui: Set up the Streamlit web interface with custom HTML elements.
|
|
sidebar: Configure the Streamlit sidebar for model and inference settings.
|
|
source_upload: Handle video file uploads through the Streamlit interface.
|
|
configure: Configure the model and load selected classes for inference.
|
|
inference: Perform real-time object detection inference.
|
|
|
|
Examples:
|
|
Create an Inference instance with a custom model
|
|
>>> inf = Inference(model="path/to/model.pt")
|
|
>>> inf.inference()
|
|
|
|
Create an Inference instance with default settings
|
|
>>> inf = Inference()
|
|
>>> inf.inference()
|
|
"""
|
|
|
|
def __init__(self, **kwargs: Any) -> None:
|
|
"""
|
|
Initialize the Inference class, checking Streamlit requirements and setting up the model path.
|
|
|
|
Args:
|
|
**kwargs (Any): Additional keyword arguments for model configuration.
|
|
"""
|
|
check_requirements("streamlit>=1.29.0") # scope imports for faster ultralytics package load speeds
|
|
import streamlit as st
|
|
|
|
self.st = st # Reference to the Streamlit module
|
|
self.source = None # Video source selection (webcam or video file)
|
|
self.enable_trk = False # Flag to toggle object tracking
|
|
self.conf = 0.25 # Confidence threshold for detection
|
|
self.iou = 0.45 # Intersection-over-Union (IoU) threshold for non-maximum suppression
|
|
self.org_frame = None # Container for the original frame display
|
|
self.ann_frame = None # Container for the annotated frame display
|
|
self.vid_file_name = None # Video file name or webcam index
|
|
self.selected_ind: List[int] = [] # List of selected class indices for detection
|
|
self.model = None # YOLO model instance
|
|
|
|
self.temp_dict = {"model": None, **kwargs}
|
|
self.model_path = None # Model file path
|
|
if self.temp_dict["model"] is not None:
|
|
self.model_path = self.temp_dict["model"]
|
|
|
|
LOGGER.info(f"Ultralytics Solutions: ✅ {self.temp_dict}")
|
|
|
|
def web_ui(self) -> None:
|
|
"""Set up the Streamlit web interface with custom HTML elements."""
|
|
menu_style_cfg = """<style>MainMenu {visibility: hidden;}</style>""" # Hide main menu style
|
|
|
|
# Main title of streamlit application
|
|
main_title_cfg = """<div><h1 style="color:#FF64DA; text-align:center; font-size:40px; margin-top:-50px;
|
|
font-family: 'Archivo', sans-serif; margin-bottom:20px;">Ultralytics YOLO Streamlit Application</h1></div>"""
|
|
|
|
# Subtitle of streamlit application
|
|
sub_title_cfg = """<div><h4 style="color:#042AFF; text-align:center; font-family: 'Archivo', sans-serif;
|
|
margin-top:-15px; margin-bottom:50px;">Experience real-time object detection on your webcam with the power
|
|
of Ultralytics YOLO! 🚀</h4></div>"""
|
|
|
|
# Set html page configuration and append custom HTML
|
|
self.st.set_page_config(page_title="Ultralytics Streamlit App", layout="wide")
|
|
self.st.markdown(menu_style_cfg, unsafe_allow_html=True)
|
|
self.st.markdown(main_title_cfg, unsafe_allow_html=True)
|
|
self.st.markdown(sub_title_cfg, unsafe_allow_html=True)
|
|
|
|
def sidebar(self) -> None:
|
|
"""Configure the Streamlit sidebar for model and inference settings."""
|
|
with self.st.sidebar: # Add Ultralytics LOGO
|
|
logo = "https://raw.githubusercontent.com/ultralytics/assets/main/logo/Ultralytics_Logotype_Original.svg"
|
|
self.st.image(logo, width=250)
|
|
|
|
self.st.sidebar.title("User Configuration") # Add elements to vertical setting menu
|
|
self.source = self.st.sidebar.selectbox(
|
|
"Video",
|
|
("webcam", "video"),
|
|
) # Add source selection dropdown
|
|
self.enable_trk = self.st.sidebar.radio("Enable Tracking", ("Yes", "No")) == "Yes" # Enable object tracking
|
|
self.conf = float(
|
|
self.st.sidebar.slider("Confidence Threshold", 0.0, 1.0, self.conf, 0.01)
|
|
) # Slider for confidence
|
|
self.iou = float(self.st.sidebar.slider("IoU Threshold", 0.0, 1.0, self.iou, 0.01)) # Slider for NMS threshold
|
|
|
|
col1, col2 = self.st.columns(2) # Create two columns for displaying frames
|
|
self.org_frame = col1.empty() # Container for original frame
|
|
self.ann_frame = col2.empty() # Container for annotated frame
|
|
|
|
def source_upload(self) -> None:
|
|
"""Handle video file uploads through the Streamlit interface."""
|
|
self.vid_file_name = ""
|
|
if self.source == "video":
|
|
vid_file = self.st.sidebar.file_uploader("Upload Video File", type=["mp4", "mov", "avi", "mkv"])
|
|
if vid_file is not None:
|
|
g = io.BytesIO(vid_file.read()) # BytesIO Object
|
|
with open("ultralytics.mp4", "wb") as out: # Open temporary file as bytes
|
|
out.write(g.read()) # Read bytes into file
|
|
self.vid_file_name = "ultralytics.mp4"
|
|
elif self.source == "webcam":
|
|
self.vid_file_name = 0 # Use webcam index 0
|
|
|
|
def configure(self) -> None:
|
|
"""Configure the model and load selected classes for inference."""
|
|
# Add dropdown menu for model selection
|
|
M_ORD, T_ORD = ["yolo11n", "yolo11s", "yolo11m", "yolo11l", "yolo11x"], ["", "-seg", "-pose", "-obb", "-cls"]
|
|
available_models = sorted(
|
|
[
|
|
x.replace("yolo", "YOLO")
|
|
for x in GITHUB_ASSETS_STEMS
|
|
if any(x.startswith(b) for b in M_ORD) and "grayscale" not in x
|
|
],
|
|
key=lambda x: (M_ORD.index(x[:7].lower()), T_ORD.index(x[7:].lower() or "")),
|
|
)
|
|
if self.model_path: # If user provided the custom model, insert model without suffix as *.pt is added later
|
|
available_models.insert(0, self.model_path.split(".pt", 1)[0])
|
|
selected_model = self.st.sidebar.selectbox("Model", available_models)
|
|
|
|
with self.st.spinner("Model is downloading..."):
|
|
self.model = YOLO(f"{selected_model.lower()}.pt") # Load the YOLO model
|
|
class_names = list(self.model.names.values()) # Convert dictionary to list of class names
|
|
self.st.success("Model loaded successfully!")
|
|
|
|
# Multiselect box with class names and get indices of selected classes
|
|
selected_classes = self.st.sidebar.multiselect("Classes", class_names, default=class_names[:3])
|
|
self.selected_ind = [class_names.index(option) for option in selected_classes]
|
|
|
|
if not isinstance(self.selected_ind, list): # Ensure selected_options is a list
|
|
self.selected_ind = list(self.selected_ind)
|
|
|
|
def inference(self) -> None:
|
|
"""Perform real-time object detection inference on video or webcam feed."""
|
|
self.web_ui() # Initialize the web interface
|
|
self.sidebar() # Create the sidebar
|
|
self.source_upload() # Upload the video source
|
|
self.configure() # Configure the app
|
|
|
|
if self.st.sidebar.button("Start"):
|
|
stop_button = self.st.button("Stop") # Button to stop the inference
|
|
cap = cv2.VideoCapture(self.vid_file_name) # Capture the video
|
|
if not cap.isOpened():
|
|
self.st.error("Could not open webcam or video source.")
|
|
return
|
|
|
|
while cap.isOpened():
|
|
success, frame = cap.read()
|
|
if not success:
|
|
self.st.warning("Failed to read frame from webcam. Please verify the webcam is connected properly.")
|
|
break
|
|
|
|
# Process frame with model
|
|
if self.enable_trk:
|
|
results = self.model.track(
|
|
frame, conf=self.conf, iou=self.iou, classes=self.selected_ind, persist=True
|
|
)
|
|
else:
|
|
results = self.model(frame, conf=self.conf, iou=self.iou, classes=self.selected_ind)
|
|
|
|
annotated_frame = results[0].plot() # Add annotations on frame
|
|
|
|
if stop_button:
|
|
cap.release() # Release the capture
|
|
self.st.stop() # Stop streamlit app
|
|
|
|
self.org_frame.image(frame, channels="BGR") # Display original frame
|
|
self.ann_frame.image(annotated_frame, channels="BGR") # Display processed frame
|
|
|
|
cap.release() # Release the capture
|
|
cv2.destroyAllWindows() # Destroy all OpenCV windows
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys # Import the sys module for accessing command-line arguments
|
|
|
|
# Check if a model name is provided as a command-line argument
|
|
args = len(sys.argv)
|
|
model = sys.argv[1] if args > 1 else None # Assign first argument as the model name if provided
|
|
# Create an instance of the Inference class and run inference
|
|
Inference(model=model).inference()
|