Reference for ultralytics/solutions/streamlit_inference.py#

class Inference:
    """A class to perform Ultralytics YOLO inference in a Streamlit web application.

    This class provides functionalities for loading models, configuring settings, uploading video or image files, and
    performing real-time inference using Streamlit and Ultralytics YOLO models for tasks such as detection,
    segmentation, semantic segmentation, depth estimation, classification, pose estimation, and oriented bounding boxes.

    Attributes:
        st (module): Streamlit module for UI creation.
        model_path (str | None): Custom model path added to the top of the model selection list.
        model (YOLO): The YOLO model instance.
        source (str): Selected input source ("webcam", "video", or "image").
        img_file_names (list[dict[str, str]]): Uploaded images as dicts with temporary file "path" and original "name".
        imgsz (int): Inference image size.
        enable_trk (bool): Enable tracking option.
        conf (float): Confidence threshold for detection, adjustable in the sidebar.
        iou (float): IoU threshold for non-maximum suppression, adjustable in the sidebar.
        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 and image file uploads through the Streamlit interface.
        configure: Configure the model and load selected classes for inference.
        image_inference: Perform inference on uploaded images.
        inference: Run the Streamlit app and perform inference on the selected source.

    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, model: str | None = None, imgsz: int = 640, conf: float = 0.25, iou: float = 0.7) -> None:
        """Initialize the Inference class, checking Streamlit requirements and setting up the model path.

        Args:
            model (str, optional): Custom model path offered first in the model selection list.
            imgsz (int): Inference image size.
            conf (float): Initial confidence threshold for the sidebar slider.
            iou (float): Initial NMS IoU threshold for the sidebar slider.
        """
        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.img_file_names = []  # List of image file names
        self.enable_trk = False  # Flag to toggle object tracking
        self.conf = conf  # Confidence threshold for detection
        self.iou = iou  # 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.model_path = model  # Custom model file path
        self.imgsz = imgsz

        LOGGER.info(f"Ultralytics Solutions: ✅ model={model}, imgsz={imgsz}, conf={conf}, iou={iou}")

def configure(self) -> None:
    """Configure the model and load selected classes for inference."""
    # Add dropdown menu for model selection, offering the standard size/task grid in order
    available_models = [
        f"{size}{task}".replace("yolo", "YOLO")
        for size in ("yolo26n", "yolo26s", "yolo26m", "yolo26l", "yolo26x")
        for task in ("", "-seg", "-sem", "-depth", "-pose", "-obb", "-cls")
        if f"{size}{task}" in GITHUB_ASSETS_STEMS
    ]
    if self.model_path:  # Insert user provided custom model in available_models
        available_models.insert(0, self.model_path)
    selected_model = self.st.sidebar.selectbox("Model", available_models)

    with self.st.spinner("Model is downloading..."):
        if selected_model.endswith((".pt", ".onnx", ".torchscript", ".mlpackage", ".engine")) or any(
            fmt in selected_model for fmt in ("openvino_model", "rknn_model")
        ):
            model_path = selected_model
        else:
            model_path = f"{selected_model.lower()}.pt"  # Default to .pt if no model provided during function call.
        self.model = YOLO(model_path)  # 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 image_inference(self) -> None:
    """Perform inference on uploaded images."""
    for img_info in self.img_file_names:
        img_path = img_info["path"]
        image = cv2.imread(img_path)  # Load and display the original image
        if image is not None:
            self.st.markdown(f"#### Processed: {img_info['name']}")
            col1, col2 = self.st.columns(2)
            with col1:
                self.st.image(image, channels="BGR", caption="Original Image")
            results = self.model(image, conf=self.conf, iou=self.iou, classes=self.selected_ind, imgsz=self.imgsz)
            annotated_image = results[0].plot()
            with col2:
                self.st.image(annotated_image, channels="BGR", caption="Predicted Image")
            try:  # Clean up temporary file
                os.unlink(img_path)
            except FileNotFoundError:
                pass  # File doesn't exist, ignore
        else:
            self.st.error("Could not load the uploaded image.")

def inference(self) -> None:
    """Run the Streamlit app and perform inference on the selected webcam, video, or image source."""
    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"):
        if self.source == "image":
            if self.img_file_names:
                self.image_inference()
            else:
                self.st.info("Please upload an image file to perform inference.")
            return

        stop_button = self.st.sidebar.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, imgsz=self.imgsz, persist=True
                )
            else:
                results = self.model(
                    frame, conf=self.conf, iou=self.iou, classes=self.selected_ind, imgsz=self.imgsz
                )

            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", caption="Original Frame")  # Display original frame
            self.ann_frame.image(annotated_frame, channels="BGR", caption="Predicted Frame")  # Display processed

        cap.release()  # Release the capture

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(
        "Source",
        ("webcam", "video", "image"),
    )  # Add source selection dropdown
    if self.source in ["webcam", "video"]:
        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

    if self.source != "image":  # Only create columns for video/webcam
        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 and image file uploads, or webcam selection, through the Streamlit interface."""
    from ultralytics.data.utils import IMG_FORMATS, VID_FORMATS  # scope import

    self.vid_file_name = ""
    if self.source == "video":
        vid_file = self.st.sidebar.file_uploader("Upload Video File", type=VID_FORMATS)
        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
    elif self.source == "image":
        import tempfile  # scope import

        if imgfiles := self.st.sidebar.file_uploader(
            "Upload Image Files", type=IMG_FORMATS, accept_multiple_files=True
        ):
            for imgfile in imgfiles:  # Save each uploaded image to a temporary file
                with tempfile.NamedTemporaryFile(delete=False, suffix=f".{imgfile.name.split('.')[-1]}") as tf:
                    tf.write(imgfile.read())
                    self.img_file_names.append({"path": tf.name, "name": imgfile.name})

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:#111F68; 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><h5 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, videos, and images
    with the power of Ultralytics YOLO! 🚀</h5></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)