ultralytics 8.1.29__py3-none-any.whl → 8.3.63__py3-none-any.whl

ultralytics/solutions/solutions.py

@@ -0,0 +1,178 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from collections import defaultdict
+
+import cv2
+
+from ultralytics import YOLO
+from ultralytics.utils import ASSETS_URL, DEFAULT_CFG_DICT, DEFAULT_SOL_DICT, LOGGER
+from ultralytics.utils.checks import check_imshow, check_requirements
+
+
+class BaseSolution:
+    """
+    A base class for managing Ultralytics Solutions.
+
+    This class provides core functionality for various Ultralytics Solutions, including model loading, object tracking,
+    and region initialization.
+
+    Attributes:
+        LineString (shapely.geometry.LineString): Class for creating line string geometries.
+        Polygon (shapely.geometry.Polygon): Class for creating polygon geometries.
+        Point (shapely.geometry.Point): Class for creating point geometries.
+        CFG (Dict): Configuration dictionary loaded from a YAML file and updated with kwargs.
+        region (List[Tuple[int, int]]): List of coordinate tuples defining a region of interest.
+        line_width (int): Width of lines used in visualizations.
+        model (ultralytics.YOLO): Loaded YOLO model instance.
+        names (Dict[int, str]): Dictionary mapping class indices to class names.
+        env_check (bool): Flag indicating whether the environment supports image display.
+        track_history (collections.defaultdict): Dictionary to store tracking history for each object.
+
+    Methods:
+        extract_tracks: Apply object tracking and extract tracks from an input image.
+        store_tracking_history: Store object tracking history for a given track ID and bounding box.
+        initialize_region: Initialize the counting region and line segment based on configuration.
+        display_output: Display the results of processing, including showing frames or saving results.
+
+    Examples:
+        >>> solution = BaseSolution(model="yolov8n.pt", region=[(0, 0), (100, 0), (100, 100), (0, 100)])
+        >>> solution.initialize_region()
+        >>> image = cv2.imread("image.jpg")
+        >>> solution.extract_tracks(image)
+        >>> solution.display_output(image)
+    """
+
+    def __init__(self, IS_CLI=False, **kwargs):
+        """
+        Initializes the `BaseSolution` class with configuration settings and the YOLO model for Ultralytics solutions.
+
+        IS_CLI (optional): Enables CLI mode if set.
+        """
+        check_requirements("shapely>=2.0.0")
+        from shapely.geometry import LineString, Point, Polygon
+        from shapely.prepared import prep
+
+        self.LineString = LineString
+        self.Polygon = Polygon
+        self.Point = Point
+        self.prep = prep
+        self.annotator = None  # Initialize annotator
+        self.tracks = None
+        self.track_data = None
+        self.boxes = []
+        self.clss = []
+        self.track_ids = []
+        self.track_line = None
+        self.r_s = None
+
+        # Load config and update with args
+        DEFAULT_SOL_DICT.update(kwargs)
+        DEFAULT_CFG_DICT.update(kwargs)
+        self.CFG = {**DEFAULT_SOL_DICT, **DEFAULT_CFG_DICT}
+        LOGGER.info(f"Ultralytics Solutions: ✅ {DEFAULT_SOL_DICT}")
+
+        self.region = self.CFG["region"]  # Store region data for other classes usage
+        self.line_width = (
+            self.CFG["line_width"] if self.CFG["line_width"] is not None else 2
+        )  # Store line_width for usage
+
+        # Load Model and store classes names
+        if self.CFG["model"] is None:
+            self.CFG["model"] = "yolo11n.pt"
+        self.model = YOLO(self.CFG["model"])
+        self.names = self.model.names
+
+        self.track_add_args = {  # Tracker additional arguments for advance configuration
+            k: self.CFG[k] for k in ["verbose", "iou", "conf", "device", "max_det", "half", "tracker"]
+        }
+
+        if IS_CLI and self.CFG["source"] is None:
+            d_s = "solutions_ci_demo.mp4" if "-pose" not in self.CFG["model"] else "solution_ci_pose_demo.mp4"
+            LOGGER.warning(f"⚠️ WARNING: source not provided. using default source {ASSETS_URL}/{d_s}")
+            from ultralytics.utils.downloads import safe_download
+
+            safe_download(f"{ASSETS_URL}/{d_s}")  # download source from ultralytics assets
+            self.CFG["source"] = d_s  # set default source
+
+        # Initialize environment and region setup
+        self.env_check = check_imshow(warn=True)
+        self.track_history = defaultdict(list)
+
+    def extract_tracks(self, im0):
+        """
+        Applies object tracking and extracts tracks from an input image or frame.
+
+        Args:
+            im0 (ndarray): The input image or frame.
+
+        Examples:
+            >>> solution = BaseSolution()
+            >>> frame = cv2.imread("path/to/image.jpg")
+            >>> solution.extract_tracks(frame)
+        """
+        self.tracks = self.model.track(source=im0, persist=True, classes=self.CFG["classes"], **self.track_add_args)
+
+        # Extract tracks for OBB or object detection
+        self.track_data = self.tracks[0].obb or self.tracks[0].boxes
+
+        if self.track_data and self.track_data.id is not None:
+            self.boxes = self.track_data.xyxy.cpu()
+            self.clss = self.track_data.cls.cpu().tolist()
+            self.track_ids = self.track_data.id.int().cpu().tolist()
+        else:
+            LOGGER.warning("WARNING ⚠️ no tracks found!")
+            self.boxes, self.clss, self.track_ids = [], [], []
+
+    def store_tracking_history(self, track_id, box):
+        """
+        Stores the tracking history of an object.
+
+        This method updates the tracking history for a given object by appending the center point of its
+        bounding box to the track line. It maintains a maximum of 30 points in the tracking history.
+
+        Args:
+            track_id (int): The unique identifier for the tracked object.
+            box (List[float]): The bounding box coordinates of the object in the format [x1, y1, x2, y2].
+
+        Examples:
+            >>> solution = BaseSolution()
+            >>> solution.store_tracking_history(1, [100, 200, 300, 400])
+        """
+        # Store tracking history
+        self.track_line = self.track_history[track_id]
+        self.track_line.append(((box[0] + box[2]) / 2, (box[1] + box[3]) / 2))
+        if len(self.track_line) > 30:
+            self.track_line.pop(0)
+
+    def initialize_region(self):
+        """Initialize the counting region and line segment based on configuration settings."""
+        if self.region is None:
+            self.region = [(20, 400), (1080, 400), (1080, 360), (20, 360)]
+        self.r_s = (
+            self.Polygon(self.region) if len(self.region) >= 3 else self.LineString(self.region)
+        )  # region or line
+
+    def display_output(self, im0):
+        """
+        Display the results of the processing, which could involve showing frames, printing counts, or saving results.
+
+        This method is responsible for visualizing the output of the object detection and tracking process. It displays
+        the processed frame with annotations, and allows for user interaction to close the display.
+
+        Args:
+            im0 (numpy.ndarray): The input image or frame that has been processed and annotated.
+
+        Examples:
+            >>> solution = BaseSolution()
+            >>> frame = cv2.imread("path/to/image.jpg")
+            >>> solution.display_output(frame)
+
+        Notes:
+            - This method will only display output if the 'show' configuration is set to True and the environment
+              supports image display.
+            - The display can be closed by pressing the 'q' key.
+        """
+        if self.CFG.get("show") and self.env_check:
+            cv2.imshow("Ultralytics Solutions", im0)
+            if cv2.waitKey(1) & 0xFF == ord("q"):
+                return

ultralytics/solutions/speed_estimation.py

@@ -1,198 +1,110 @@
-# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from collections import defaultdict
 from time import time
 
-import cv2
 import numpy as np
 
-from ultralytics.utils.checks import check_imshow
+from ultralytics.solutions.solutions import BaseSolution
 from ultralytics.utils.plotting import Annotator, colors
 
 
-class SpeedEstimator:
-    """A class to estimation speed of objects in real-time video stream based on their tracks."""
-
-    def __init__(self):
-        """Initializes the speed-estimator class with default values for Visual, Image, track and speed parameters."""
-
-        # Visual & im0 information
-        self.im0 = None
-        self.annotator = None
-        self.view_img = False
-
-        # Region information
-        self.reg_pts = [(20, 400), (1260, 400)]
-        self.region_thickness = 3
-
-        # Predict/track information
-        self.clss = None
-        self.names = None
-        self.boxes = None
-        self.trk_ids = None
-        self.trk_pts = None
-        self.line_thickness = 2
-        self.trk_history = defaultdict(list)
-
-        # Speed estimator information
-        self.current_time = 0
-        self.dist_data = {}
-        self.trk_idslist = []
-        self.spdl_dist_thresh = 10
-        self.trk_previous_times = {}
-        self.trk_previous_points = {}
-
-        # Check if environment support imshow
-        self.env_check = check_imshow(warn=True)
-
-    def set_args(
-        self,
-        reg_pts,
-        names,
-        view_img=False,
-        line_thickness=2,
-        region_thickness=5,
-        spdl_dist_thresh=10,
-    ):
+class SpeedEstimator(BaseSolution):
+    """
+    A class to estimate the speed of objects in a real-time video stream based on their tracks.
+
+    This class extends the BaseSolution class and provides functionality for estimating object speeds using
+    tracking data in video streams.
+
+    Attributes:
+        spd (Dict[int, float]): Dictionary storing speed data for tracked objects.
+        trkd_ids (List[int]): List of tracked object IDs that have already been speed-estimated.
+        trk_pt (Dict[int, float]): Dictionary storing previous timestamps for tracked objects.
+        trk_pp (Dict[int, Tuple[float, float]]): Dictionary storing previous positions for tracked objects.
+        annotator (Annotator): Annotator object for drawing on images.
+        region (List[Tuple[int, int]]): List of points defining the speed estimation region.
+        track_line (List[Tuple[float, float]]): List of points representing the object's track.
+        r_s (LineString): LineString object representing the speed estimation region.
+
+    Methods:
+        initialize_region: Initializes the speed estimation region.
+        estimate_speed: Estimates the speed of objects based on tracking data.
+        store_tracking_history: Stores the tracking history for an object.
+        extract_tracks: Extracts tracks from the current frame.
+        display_output: Displays the output with annotations.
+
+    Examples:
+        >>> estimator = SpeedEstimator()
+        >>> frame = cv2.imread("frame.jpg")
+        >>> processed_frame = estimator.estimate_speed(frame)
+        >>> cv2.imshow("Speed Estimation", processed_frame)
+    """
+
+    def __init__(self, **kwargs):
+        """Initializes the SpeedEstimator object with speed estimation parameters and data structures."""
+        super().__init__(**kwargs)
+
+        self.initialize_region()  # Initialize speed region
+
+        self.spd = {}  # set for speed data
+        self.trkd_ids = []  # list for already speed_estimated and tracked ID's
+        self.trk_pt = {}  # set for tracks previous time
+        self.trk_pp = {}  # set for tracks previous point
+
+    def estimate_speed(self, im0):
         """
-        Configures the speed estimation and display parameters.
+        Estimates the speed of objects based on tracking data.
 
         Args:
-            reg_pts (list): Initial list of points defining the speed calculation region.
-            names (dict): object detection classes names
-            view_img (bool): Flag indicating frame display
-            line_thickness (int): Line thickness for bounding boxes.
-            region_thickness (int): Speed estimation region thickness
-            spdl_dist_thresh (int): Euclidean distance threshold for speed line
-        """
-        if reg_pts is None:
-            print("Region points not provided, using default values")
-        else:
-            self.reg_pts = reg_pts
-        self.names = names
-        self.view_img = view_img
-        self.line_thickness = line_thickness
-        self.region_thickness = region_thickness
-        self.spdl_dist_thresh = spdl_dist_thresh
-
-    def extract_tracks(self, tracks):
-        """
-        Extracts results from the provided data.
-
-        Args:
-            tracks (list): List of tracks obtained from the object tracking process.
-        """
-        self.boxes = tracks[0].boxes.xyxy.cpu()
-        self.clss = tracks[0].boxes.cls.cpu().tolist()
-        self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()
-
-    def store_track_info(self, track_id, box):
-        """
-        Store track data.
-
-        Args:
-            track_id (int): object track id.
-            box (list): object bounding box data
-        """
-        track = self.trk_history[track_id]
-        bbox_center = (float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2))
-        track.append(bbox_center)
-
-        if len(track) > 30:
-            track.pop(0)
-
-        self.trk_pts = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
-        return track
+            im0 (np.ndarray): Input image for processing. Shape is typically (H, W, C) for RGB images.
 
-    def plot_box_and_track(self, track_id, box, cls, track):
-        """
-        Plot track and bounding box.
+        Returns:
+            (np.ndarray): Processed image with speed estimations and annotations.
 
-        Args:
-            track_id (int): object track id.
-            box (list): object bounding box data
-            cls (str): object class name
-            track (list): tracking history for tracks path drawing
+        Examples:
+            >>> estimator = SpeedEstimator()
+            >>> image = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
+            >>> processed_image = estimator.estimate_speed(image)
         """
-        speed_label = f"{int(self.dist_data[track_id])}km/ph" if track_id in self.dist_data else self.names[int(cls)]
-        bbox_color = colors(int(track_id)) if track_id in self.dist_data else (255, 0, 255)
-
-        self.annotator.box_label(box, speed_label, bbox_color)
-
-        cv2.polylines(self.im0, [self.trk_pts], isClosed=False, color=(0, 255, 0), thickness=1)
-        cv2.circle(self.im0, (int(track[-1][0]), int(track[-1][1])), 5, bbox_color, -1)
+        self.annotator = Annotator(im0, line_width=self.line_width)  # Initialize annotator
+        self.extract_tracks(im0)  # Extract tracks
 
-    def calculate_speed(self, trk_id, track):
-        """
-        Calculation of object speed.
-
-        Args:
-            trk_id (int): object track id.
-            track (list): tracking history for tracks path drawing
-        """
-
-        if not self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]:
-            return
-        if self.reg_pts[1][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[1][1] + self.spdl_dist_thresh:
-            direction = "known"
-
-        elif self.reg_pts[0][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[0][1] + self.spdl_dist_thresh:
-            direction = "known"
-
-        else:
-            direction = "unknown"
-
-        if self.trk_previous_times[trk_id] != 0 and direction != "unknown" and trk_id not in self.trk_idslist:
-            self.trk_idslist.append(trk_id)
-
-            time_difference = time() - self.trk_previous_times[trk_id]
-            if time_difference > 0:
-                dist_difference = np.abs(track[-1][1] - self.trk_previous_points[trk_id][1])
-                speed = dist_difference / time_difference
-                self.dist_data[trk_id] = speed
-
-        self.trk_previous_times[trk_id] = time()
-        self.trk_previous_points[trk_id] = track[-1]
-
-    def estimate_speed(self, im0, tracks, region_color=(255, 0, 0)):
-        """
-        Calculate object based on tracking data.
-
-        Args:
-            im0 (nd array): Image
-            tracks (list): List of tracks obtained from the object tracking process.
-            region_color (tuple): Color to use when drawing regions.
-        """
-        self.im0 = im0
-        if tracks[0].boxes.id is None:
-            if self.view_img and self.env_check:
-                self.display_frames()
-            return im0
-        self.extract_tracks(tracks)
+        self.annotator.draw_region(
+            reg_pts=self.region, color=(104, 0, 123), thickness=self.line_width * 2
+        )  # Draw region
 
-        self.annotator = Annotator(self.im0, line_width=2)
-        self.annotator.draw_region(reg_pts=self.reg_pts, color=region_color, thickness=self.region_thickness)
+        for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
+            self.store_tracking_history(track_id, box)  # Store track history
 
-        for box, trk_id, cls in zip(self.boxes, self.trk_ids, self.clss):
-            track = self.store_track_info(trk_id, box)
+            # Check if track_id is already in self.trk_pp or trk_pt initialize if not
+            if track_id not in self.trk_pt:
+                self.trk_pt[track_id] = 0
+            if track_id not in self.trk_pp:
+                self.trk_pp[track_id] = self.track_line[-1]
 
-            if trk_id not in self.trk_previous_times:
-                self.trk_previous_times[trk_id] = 0
+            speed_label = f"{int(self.spd[track_id])} km/h" if track_id in self.spd else self.names[int(cls)]
+            self.annotator.box_label(box, label=speed_label, color=colors(track_id, True))  # Draw bounding box
 
-            self.plot_box_and_track(trk_id, box, cls, track)
-            self.calculate_speed(trk_id, track)
+            # Draw tracks of objects
+            self.annotator.draw_centroid_and_tracks(
+                self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width
+            )
 
-        if self.view_img and self.env_check:
-            self.display_frames()
+            # Calculate object speed and direction based on region intersection
+            if self.LineString([self.trk_pp[track_id], self.track_line[-1]]).intersects(self.r_s):
+                direction = "known"
+            else:
+                direction = "unknown"
 
-        return im0
+            # Perform speed calculation and tracking updates if direction is valid
+            if direction == "known" and track_id not in self.trkd_ids:
+                self.trkd_ids.append(track_id)
+                time_difference = time() - self.trk_pt[track_id]
+                if time_difference > 0:
+                    self.spd[track_id] = np.abs(self.track_line[-1][1] - self.trk_pp[track_id][1]) / time_difference
 
-    def display_frames(self):
-        """Display frame."""
-        cv2.imshow("Ultralytics Speed Estimation", self.im0)
-        if cv2.waitKey(1) & 0xFF == ord("q"):
-            return
+            self.trk_pt[track_id] = time()
+            self.trk_pp[track_id] = self.track_line[-1]
 
+        self.display_output(im0)  # display output with base class function
 
-if __name__ == "__main__":
-    SpeedEstimator()
+        return im0  # return output image for more usage

ultralytics/solutions/streamlit_inference.py

@@ -0,0 +1,190 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import io
+from typing import Any
+
+import cv2
+
+from ultralytics import YOLO
+from ultralytics.utils import LOGGER
+from ultralytics.utils.checks import check_requirements
+from ultralytics.utils.downloads import GITHUB_ASSETS_STEMS
+
+
+class Inference:
+    """
+    A class to perform object detection, image classification, image segmentation and pose estimation inference using
+    Streamlit and Ultralytics YOLO models. It provides the functionalities such as loading models, configuring settings,
+    uploading video files, and performing real-time inference.
+
+    Attributes:
+        st (module): Streamlit module for UI creation.
+        temp_dict (dict): Temporary dictionary to store the model path.
+        model_path (str): Path to the loaded model.
+        model (YOLO): The YOLO model instance.
+        source (str): Selected video source.
+        enable_trk (str): Enable tracking option.
+        conf (float): Confidence threshold.
+        iou (float): IoU threshold for non-max suppression.
+        vid_file_name (str): Name of the uploaded video file.
+        selected_ind (list): List of selected class indices.
+
+    Methods:
+        web_ui: Sets up the Streamlit web interface with custom HTML elements.
+        sidebar: Configures the Streamlit sidebar for model and inference settings.
+        source_upload: Handles video file uploads through the Streamlit interface.
+        configure: Configures the model and loads selected classes for inference.
+        inference: Performs real-time object detection inference.
+
+    Examples:
+        >>> inf = solutions.Inference(model="path/to/model.pt")  # Model is not necessary argument.
+        >>> inf.inference()
+    """
+
+    def __init__(self, **kwargs: Any):
+        """
+        Initializes 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 class instance
+        self.source = None  # Placeholder for video or webcam source details
+        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 to be displayed
+        self.ann_frame = None  # Container for the annotated frame to be displayed
+        self.vid_file_name = None  # Holds the name of the video file
+        self.selected_ind = []  # List of selected classes for detection or tracking
+        self.model = None  # Container for the loaded model instance
+
+        self.temp_dict = {"model": None, **kwargs}
+        self.model_path = None  # Store model file name with 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):
+        """Sets 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):
+        """Configures 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"))  # 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)
+        self.org_frame = col1.empty()
+        self.ann_frame = col2.empty()
+
+    def source_upload(self):
+        """Handles 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
+
+    def configure(self):
+        """Configures the model and loads selected classes for inference."""
+        # Add dropdown menu for model selection
+        available_models = [x.replace("yolo", "YOLO") for x in GITHUB_ASSETS_STEMS if x.startswith("yolo11")]
+        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")[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):
+        """Performs real-time object detection inference."""
+        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.")
+            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
+
+                # Store model predictions
+                if self.enable_trk == "Yes":
+                    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 window
+
+
+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
+    # Create an instance of the Inference class and run inference
+    Inference(model=model).inference()