ultralytics 8.3.89__py3-none-any.whl → 8.3.90__py3-none-any.whl

ultralytics/trackers/utils/gmc.py

@@ -16,20 +16,19 @@ class GMC:
     SIFT, ECC, and Sparse Optical Flow. It also supports downscaling of frames for computational efficiency.
 
     Attributes:
-        method (str): The method used for tracking. Options include 'orb', 'sift', 'ecc', 'sparseOptFlow', 'none'.
+        method (str): The tracking method to use. Options include 'orb', 'sift', 'ecc', 'sparseOptFlow', 'none'.
         downscale (int): Factor by which to downscale the frames for processing.
-        prevFrame (np.ndarray): Stores the previous frame for tracking.
-        prevKeyPoints (List): Stores the keypoints from the previous frame.
-        prevDescriptors (np.ndarray): Stores the descriptors from the previous frame.
-        initializedFirstFrame (bool): Flag to indicate if the first frame has been processed.
+        prevFrame (np.ndarray): Previous frame for tracking.
+        prevKeyPoints (List): Keypoints from the previous frame.
+        prevDescriptors (np.ndarray): Descriptors from the previous frame.
+        initializedFirstFrame (bool): Flag indicating if the first frame has been processed.
 
     Methods:
-        __init__: Initializes a GMC object with the specified method and downscale factor.
-        apply: Applies the chosen method to a raw frame and optionally uses provided detections.
-        apply_ecc: Applies the ECC algorithm to a raw frame.
-        apply_features: Applies feature-based methods like ORB or SIFT to a raw frame.
-        apply_sparseoptflow: Applies the Sparse Optical Flow method to a raw frame.
-        reset_params: Resets the internal parameters of the GMC object.
+        apply: Apply the chosen method to a raw frame and optionally use provided detections.
+        apply_ecc: Apply the ECC algorithm to a raw frame.
+        apply_features: Apply feature-based methods like ORB or SIFT to a raw frame.
+        apply_sparseoptflow: Apply the Sparse Optical Flow method to a raw frame.
+        reset_params: Reset the internal parameters of the GMC object.
 
     Examples:
         Create a GMC object and apply it to a frame
@@ -46,7 +45,7 @@ class GMC:
         Initialize a Generalized Motion Compensation (GMC) object with tracking method and downscale factor.
 
         Args:
-            method (str): The method used for tracking. Options include 'orb', 'sift', 'ecc', 'sparseOptFlow', 'none'.
+            method (str): The tracking method to use. Options include 'orb', 'sift', 'ecc', 'sparseOptFlow', 'none'.
             downscale (int): Downscale factor for processing frames.
 
         Examples:
@@ -82,14 +81,14 @@ class GMC:
         elif self.method in {"none", "None", None}:
             self.method = None
         else:
-            raise ValueError(f"Error: Unknown GMC method:{method}")
+            raise ValueError(f"Error: Unknown GMC method: {method}")
 
         self.prevFrame = None
         self.prevKeyPoints = None
         self.prevDescriptors = None
         self.initializedFirstFrame = False
 
-    def apply(self, raw_frame: np.array, detections: list = None) -> np.array:
+    def apply(self, raw_frame: np.ndarray, detections: list = None) -> np.ndarray:
         """
         Apply object detection on a raw frame using the specified method.
 
@@ -98,14 +97,14 @@ class GMC:
             detections (List | None): List of detections to be used in the processing.
 
         Returns:
-            (np.ndarray): Processed frame with applied object detection.
+            (np.ndarray): Transformation matrix with shape (2, 3).
 
         Examples:
             >>> gmc = GMC(method="sparseOptFlow")
             >>> raw_frame = np.random.rand(480, 640, 3)
-            >>> processed_frame = gmc.apply(raw_frame)
-            >>> print(processed_frame.shape)
-            (480, 640, 3)
+            >>> transformation_matrix = gmc.apply(raw_frame)
+            >>> print(transformation_matrix.shape)
+            (2, 3)
         """
         if self.method in {"orb", "sift"}:
             return self.apply_features(raw_frame, detections)
@@ -116,7 +115,7 @@ class GMC:
         else:
             return np.eye(2, 3)
 
-    def apply_ecc(self, raw_frame: np.array) -> np.array:
+    def apply_ecc(self, raw_frame: np.ndarray) -> np.ndarray:
         """
         Apply the ECC (Enhanced Correlation Coefficient) algorithm to a raw frame for motion compensation.
 
@@ -124,7 +123,7 @@ class GMC:
             raw_frame (np.ndarray): The raw frame to be processed, with shape (H, W, C).
 
         Returns:
-            (np.ndarray): The processed frame with the applied ECC transformation.
+            (np.ndarray): Transformation matrix with shape (2, 3).
 
         Examples:
             >>> gmc = GMC(method="ecc")
@@ -161,7 +160,7 @@ class GMC:
 
         return H
 
-    def apply_features(self, raw_frame: np.array, detections: list = None) -> np.array:
+    def apply_features(self, raw_frame: np.ndarray, detections: list = None) -> np.ndarray:
         """
         Apply feature-based methods like ORB or SIFT to a raw frame.
 
@@ -170,13 +169,13 @@ class GMC:
             detections (List | None): List of detections to be used in the processing.
 
         Returns:
-            (np.ndarray): Processed frame.
+            (np.ndarray): Transformation matrix with shape (2, 3).
 
         Examples:
             >>> gmc = GMC(method="orb")
             >>> raw_frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
-            >>> processed_frame = gmc.apply_features(raw_frame)
-            >>> print(processed_frame.shape)
+            >>> transformation_matrix = gmc.apply_features(raw_frame)
+            >>> print(transformation_matrix.shape)
             (2, 3)
         """
         height, width, _ = raw_frame.shape
@@ -304,7 +303,7 @@ class GMC:
 
         return H
 
-    def apply_sparseoptflow(self, raw_frame: np.array) -> np.array:
+    def apply_sparseoptflow(self, raw_frame: np.ndarray) -> np.ndarray:
         """
         Apply Sparse Optical Flow method to a raw frame.
 
@@ -312,7 +311,7 @@ class GMC:
             raw_frame (np.ndarray): The raw frame to be processed, with shape (H, W, C).
 
         Returns:
-            (np.ndarray): Processed frame with shape (2, 3).
+            (np.ndarray): Transformation matrix with shape (2, 3).
 
         Examples:
             >>> gmc = GMC()
@@ -355,7 +354,7 @@ class GMC:
         currPoints = np.array(currPoints)
 
         # Find rigid matrix
-        if (prevPoints.shape[0] > 4) and (prevPoints.shape[0] == prevPoints.shape[0]):
+        if (prevPoints.shape[0] > 4) and (prevPoints.shape[0] == currPoints.shape[0]):
             H, _ = cv2.estimateAffinePartial2D(prevPoints, currPoints, cv2.RANSAC)
 
             if self.downscale > 1.0:

ultralytics/trackers/utils/kalman_filter.py

@@ -57,12 +57,11 @@ class KalmanFilterXYAH:
             self._motion_mat[i, ndim + i] = dt
         self._update_mat = np.eye(ndim, 2 * ndim)
 
-        # Motion and observation uncertainty are chosen relative to the current state estimate. These weights control
-        # the amount of uncertainty in the model.
+        # Motion and observation uncertainty are chosen relative to the current state estimate
         self._std_weight_position = 1.0 / 20
         self._std_weight_velocity = 1.0 / 160
 
-    def initiate(self, measurement: np.ndarray) -> tuple:
+    def initiate(self, measurement: np.ndarray):
         """
         Create a track from an unassociated measurement.
 
@@ -71,8 +70,8 @@ class KalmanFilterXYAH:
                 and height h.
 
         Returns:
-            (tuple[ndarray, ndarray]): Returns the mean vector (8-dimensional) and covariance matrix (8x8 dimensional)
-                of the new track. Unobserved velocities are initialized to 0 mean.
+            (np.ndarray): Mean vector (8-dimensional) of the new track. Unobserved velocities are initialized to 0 mean.
+            (np.ndarray): Covariance matrix (8x8 dimensional) of the new track.
 
         Examples:
             >>> kf = KalmanFilterXYAH()
@@ -96,7 +95,7 @@ class KalmanFilterXYAH:
         covariance = np.diag(np.square(std))
         return mean, covariance
 
-    def predict(self, mean: np.ndarray, covariance: np.ndarray) -> tuple:
+    def predict(self, mean: np.ndarray, covariance: np.ndarray):
         """
         Run Kalman filter prediction step.
 
@@ -105,8 +104,8 @@ class KalmanFilterXYAH:
             covariance (np.ndarray): The 8x8-dimensional covariance matrix of the object state at the previous time step.
 
         Returns:
-            (tuple[ndarray, ndarray]): Returns the mean vector and covariance matrix of the predicted state. Unobserved
-                velocities are initialized to 0 mean.
+            (np.ndarray): Mean vector of the predicted state. Unobserved velocities are initialized to 0 mean.
+            (np.ndarray): Covariance matrix of the predicted state.
 
         Examples:
             >>> kf = KalmanFilterXYAH()
@@ -133,7 +132,7 @@ class KalmanFilterXYAH:
 
         return mean, covariance
 
-    def project(self, mean: np.ndarray, covariance: np.ndarray) -> tuple:
+    def project(self, mean: np.ndarray, covariance: np.ndarray):
         """
         Project state distribution to measurement space.
 
@@ -142,7 +141,8 @@ class KalmanFilterXYAH:
             covariance (np.ndarray): The state's covariance matrix (8x8 dimensional).
 
         Returns:
-            (tuple[ndarray, ndarray]): Returns the projected mean and covariance matrix of the given state estimate.
+            (np.ndarray): Projected mean of the given state estimate.
+            (np.ndarray): Projected covariance matrix of the given state estimate.
 
         Examples:
             >>> kf = KalmanFilterXYAH()
@@ -162,7 +162,7 @@ class KalmanFilterXYAH:
         covariance = np.linalg.multi_dot((self._update_mat, covariance, self._update_mat.T))
         return mean, covariance + innovation_cov
 
-    def multi_predict(self, mean: np.ndarray, covariance: np.ndarray) -> tuple:
+    def multi_predict(self, mean: np.ndarray, covariance: np.ndarray):
         """
         Run Kalman filter prediction step for multiple object states (Vectorized version).
 
@@ -171,9 +171,8 @@ class KalmanFilterXYAH:
             covariance (np.ndarray): The Nx8x8 covariance matrix of the object states at the previous time step.
 
         Returns:
-            (tuple[ndarray, ndarray]): Returns the mean matrix and covariance matrix of the predicted states.
-                The mean matrix has shape (N, 8) and the covariance matrix has shape (N, 8, 8). Unobserved velocities
-                are initialized to 0 mean.
+            (np.ndarray): Mean matrix of the predicted states with shape (N, 8).
+            (np.ndarray): Covariance matrix of the predicted states with shape (N, 8, 8).
 
         Examples:
             >>> mean = np.random.rand(10, 8)  # 10 object states
@@ -203,7 +202,7 @@ class KalmanFilterXYAH:
 
         return mean, covariance
 
-    def update(self, mean: np.ndarray, covariance: np.ndarray, measurement: np.ndarray) -> tuple:
+    def update(self, mean: np.ndarray, covariance: np.ndarray, measurement: np.ndarray):
         """
         Run Kalman filter correction step.
 
@@ -214,7 +213,8 @@ class KalmanFilterXYAH:
                 position, a the aspect ratio, and h the height of the bounding box.
 
         Returns:
-            (tuple[ndarray, ndarray]): Returns the measurement-corrected state distribution.
+            (np.ndarray): Measurement-corrected state mean.
+            (np.ndarray): Measurement-corrected state covariance.
 
         Examples:
             >>> kf = KalmanFilterXYAH()
@@ -317,7 +317,7 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
         >>> print(covariance)
     """
 
-    def initiate(self, measurement: np.ndarray) -> tuple:
+    def initiate(self, measurement: np.ndarray):
         """
         Create track from unassociated measurement.
 
@@ -325,8 +325,8 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
             measurement (np.ndarray): Bounding box coordinates (x, y, w, h) with center position (x, y), width, and height.
 
         Returns:
-            (tuple[ndarray, ndarray]): Returns the mean vector (8 dimensional) and covariance matrix (8x8 dimensional)
-                of the new track. Unobserved velocities are initialized to 0 mean.
+            (np.ndarray): Mean vector (8 dimensional) of the new track. Unobserved velocities are initialized to 0 mean.
+            (np.ndarray): Covariance matrix (8x8 dimensional) of the new track.
 
         Examples:
             >>> kf = KalmanFilterXYWH()
@@ -361,7 +361,7 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
         covariance = np.diag(np.square(std))
         return mean, covariance
 
-    def predict(self, mean, covariance) -> tuple:
+    def predict(self, mean, covariance):
         """
         Run Kalman filter prediction step.
 
@@ -370,8 +370,8 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
             covariance (np.ndarray): The 8x8-dimensional covariance matrix of the object state at the previous time step.
 
         Returns:
-            (tuple[ndarray, ndarray]): Returns the mean vector and covariance matrix of the predicted state. Unobserved
-                velocities are initialized to 0 mean.
+            (np.ndarray): Mean vector of the predicted state. Unobserved velocities are initialized to 0 mean.
+            (np.ndarray): Covariance matrix of the predicted state.
 
         Examples:
             >>> kf = KalmanFilterXYWH()
@@ -398,7 +398,7 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
 
         return mean, covariance
 
-    def project(self, mean, covariance) -> tuple:
+    def project(self, mean, covariance):
         """
         Project state distribution to measurement space.
 
@@ -407,7 +407,8 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
             covariance (np.ndarray): The state's covariance matrix (8x8 dimensional).
 
         Returns:
-            (tuple[ndarray, ndarray]): Returns the projected mean and covariance matrix of the given state estimate.
+            (np.ndarray): Projected mean of the given state estimate.
+            (np.ndarray): Projected covariance matrix of the given state estimate.
 
         Examples:
             >>> kf = KalmanFilterXYWH()
@@ -427,7 +428,7 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
         covariance = np.linalg.multi_dot((self._update_mat, covariance, self._update_mat.T))
         return mean, covariance + innovation_cov
 
-    def multi_predict(self, mean, covariance) -> tuple:
+    def multi_predict(self, mean, covariance):
         """
         Run Kalman filter prediction step (Vectorized version).
 
@@ -436,8 +437,8 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
             covariance (np.ndarray): The Nx8x8 covariance matrix of the object states at the previous time step.
 
         Returns:
-            (tuple[ndarray, ndarray]): Returns the mean vector and covariance matrix of the predicted state. Unobserved
-                velocities are initialized to 0 mean.
+            (np.ndarray): Mean matrix of the predicted states with shape (N, 8).
+            (np.ndarray): Covariance matrix of the predicted states with shape (N, 8, 8).
 
         Examples:
             >>> mean = np.random.rand(5, 8)  # 5 objects with 8-dimensional state vectors
@@ -468,7 +469,7 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
 
         return mean, covariance
 
-    def update(self, mean, covariance, measurement) -> tuple:
+    def update(self, mean, covariance, measurement):
         """
         Run Kalman filter correction step.
 
@@ -479,7 +480,8 @@ class KalmanFilterXYWH(KalmanFilterXYAH):
                 position, w the width, and h the height of the bounding box.
 
         Returns:
-            (tuple[ndarray, ndarray]): Returns the measurement-corrected state distribution.
+            (np.ndarray): Measurement-corrected state mean.
+            (np.ndarray): Measurement-corrected state covariance.
 
         Examples:
             >>> kf = KalmanFilterXYWH()

ultralytics/trackers/utils/matching.py

@@ -66,11 +66,11 @@ def iou_distance(atracks: list, btracks: list) -> np.ndarray:
     Compute cost based on Intersection over Union (IoU) between tracks.
 
     Args:
-        atracks (list[STrack] | list[np.ndarray]): List of tracks 'a' or bounding boxes.
-        btracks (list[STrack] | list[np.ndarray]): List of tracks 'b' or bounding boxes.
+        atracks (List[STrack] | List[np.ndarray]): List of tracks 'a' or bounding boxes.
+        btracks (List[STrack] | List[np.ndarray]): List of tracks 'b' or bounding boxes.
 
     Returns:
-        (np.ndarray): Cost matrix computed based on IoU.
+        (np.ndarray): Cost matrix computed based on IoU with shape (len(atracks), len(btracks)).
 
     Examples:
         Compute IoU distance between two sets of tracks
@@ -106,8 +106,8 @@ def embedding_distance(tracks: list, detections: list, metric: str = "cosine") -
     Compute distance between tracks and detections based on embeddings.
 
     Args:
-        tracks (list[STrack]): List of tracks, where each track contains embedding features.
-        detections (list[BaseTrack]): List of detections, where each detection contains embedding features.
+        tracks (List[STrack]): List of tracks, where each track contains embedding features.
+        detections (List[BaseTrack]): List of detections, where each detection contains embedding features.
         metric (str): Metric for distance computation. Supported metrics include 'cosine', 'euclidean', etc.
 
     Returns:
@@ -133,11 +133,11 @@ def embedding_distance(tracks: list, detections: list, metric: str = "cosine") -
 
 def fuse_score(cost_matrix: np.ndarray, detections: list) -> np.ndarray:
     """
-    Fuses cost matrix with detection scores to produce a single similarity matrix.
+    Fuse cost matrix with detection scores to produce a single similarity matrix.
 
     Args:
         cost_matrix (np.ndarray): The matrix containing cost values for assignments, with shape (N, M).
-        detections (list[BaseTrack]): List of detections, each containing a score attribute.
+        detections (List[BaseTrack]): List of detections, each containing a score attribute.
 
     Returns:
         (np.ndarray): Fused similarity matrix with shape (N, M).

ultralytics/utils/__init__.py

@@ -304,17 +304,24 @@ def plt_settings(rcparams=None, backend="Agg"):
     """
     Decorator to temporarily set rc parameters and the backend for a plotting function.
 
-    Example:
-        decorator: @plt_settings({"font.size": 12})
-        context manager: with plt_settings({"font.size": 12}):
-
     Args:
-        rcparams (dict): Dictionary of rc parameters to set.
+        rcparams (dict, optional): Dictionary of rc parameters to set.
         backend (str, optional): Name of the backend to use. Defaults to 'Agg'.
 
     Returns:
-        (Callable): Decorated function with temporarily set rc parameters and backend. This decorator can be
-            applied to any function that needs to have specific matplotlib rc parameters and backend for its execution.
+        (Callable): Decorated function with temporarily set rc parameters and backend.
+
+    Examples:
+        >>> @plt_settings({"font.size": 12})
+        >>> def plot_function():
+        ...     plt.figure()
+        ...     plt.plot([1, 2, 3])
+        ...     plt.show()
+
+        >>> with plt_settings({"font.size": 12}):
+        ...     plt.figure()
+        ...     plt.plot([1, 2, 3])
+        ...     plt.show()
     """
     if rcparams is None:
         rcparams = {"font.size": 11}
@@ -357,6 +364,9 @@ def set_logging(name="LOGGING_NAME", verbose=True):
         name (str): Name of the logger. Defaults to "LOGGING_NAME".
         verbose (bool): Flag to set logging level to INFO if True, ERROR otherwise. Defaults to True.
 
+    Returns:
+        (logging.Logger): Configured logger object.
+
     Examples:
         >>> set_logging(name="ultralytics", verbose=True)
         >>> logger = logging.getLogger("ultralytics")
@@ -376,7 +386,7 @@ def set_logging(name="LOGGING_NAME", verbose=True):
 
         class CustomFormatter(logging.Formatter):
             def format(self, record):
-                """Sets up logging with UTF-8 encoding and configurable verbosity."""
+                """Format log records with UTF-8 encoding for Windows compatibility."""
                 return emojis(super().format(record))
 
         try:
@@ -420,9 +430,10 @@ def emojis(string=""):
 
 class ThreadingLocked:
     """
-    A decorator class for ensuring thread-safe execution of a function or method. This class can be used as a decorator
-    to make sure that if the decorated function is called from multiple threads, only one thread at a time will be able
-    to execute the function.
+    A decorator class for ensuring thread-safe execution of a function or method.
+
+    This class can be used as a decorator to make sure that if the decorated function is called from multiple threads,
+    only one thread at a time will be able to execute the function.
 
     Attributes:
         lock (threading.Lock): A lock object used to manage access to the decorated function.
@@ -435,7 +446,7 @@ class ThreadingLocked:
     """
 
     def __init__(self):
-        """Initializes the decorator class for thread-safe execution of a function or method."""
+        """Initialize the decorator class with a threading lock."""
         self.lock = threading.Lock()
 
     def __call__(self, f):
@@ -536,8 +547,7 @@ DEFAULT_CFG = IterableSimpleNamespace(**DEFAULT_CFG_DICT)
 
 def read_device_model() -> str:
     """
-    Reads the device model information from the system and caches it for quick access. Used by is_jetson() and
-    is_raspberrypi().
+    Reads the device model information from the system and caches it for quick access.
 
     Returns:
         (str): Kernel release information.
@@ -619,7 +629,7 @@ def is_docker() -> bool:
 
 def is_raspberrypi() -> bool:
     """
-    Determines if the Python environment is running on a Raspberry Pi by checking the device model information.
+    Determines if the Python environment is running on a Raspberry Pi.
 
     Returns:
         (bool): True if running on a Raspberry Pi, False otherwise.
@@ -629,7 +639,7 @@ def is_raspberrypi() -> bool:
 
 def is_jetson() -> bool:
     """
-    Determines if the Python environment is running on an NVIDIA Jetson device by checking the device model information.
+    Determines if the Python environment is running on an NVIDIA Jetson device.
 
     Returns:
         (bool): True if running on an NVIDIA Jetson device, False otherwise.
@@ -709,8 +719,7 @@ def is_github_action_running() -> bool:
 
 def get_git_dir():
     """
-    Determines whether the current file is part of a git repository and if so, returns the repository root directory. If
-    the current file is not part of a git repository, returns None.
+    Determines whether the current file is part of a git repository and if so, returns the repository root directory.
 
     Returns:
         (Path | None): Git root directory if found or None if not found.
@@ -722,8 +731,7 @@ def get_git_dir():
 
 def is_git_dir():
     """
-    Determines whether the current file is part of a git repository. If the current file is not part of a git
-    repository, returns None.
+    Determines whether the current file is part of a git repository.
 
     Returns:
         (bool): True if current file is part of a git repository.
@@ -1004,8 +1012,10 @@ def threaded(func):
 
 def set_sentry():
     """
-    Initialize the Sentry SDK for error tracking and reporting. Only used if sentry_sdk package is installed and
-    sync=True in settings. Run 'yolo settings' to see and update settings.
+    Initialize the Sentry SDK for error tracking and reporting.
+
+    Only used if sentry_sdk package is installed and sync=True in settings. Run 'yolo settings' to see and update
+    settings.
 
     Conditions required to send errors (ALL conditions must be met or no errors will be reported):
         - sentry_sdk package is installed
@@ -1016,11 +1026,6 @@ def set_sentry():
         - running with rank -1 or 0
         - online environment
         - CLI used to run package (checked with 'yolo' as the name of the main CLI command)
-
-    The function also configures Sentry SDK to ignore KeyboardInterrupt and FileNotFoundError exceptions and to exclude
-    events with 'out of memory' in their exception message.
-
-    Additionally, the function sets custom tags and user information for Sentry events.
     """
     if (
         not SETTINGS["sync"]

ultralytics/utils/autobatch.py

@@ -25,7 +25,7 @@ def check_train_batch_size(model, imgsz=640, amp=True, batch=-1, max_num_obj=1):
     Returns:
         (int): Optimal batch size computed using the autobatch() function.
 
-    Note:
+    Notes:
         If 0.0 < batch < 1.0, it's used as the fraction of GPU memory to use.
         Otherwise, a default fraction of 0.6 is used.
     """
@@ -40,10 +40,10 @@ def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch, max
     Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory.
 
     Args:
-        model (torch.nn.module): YOLO model to compute batch size for.
-        imgsz (int, optional): The image size used as input for the YOLO model. Defaults to 640.
-        fraction (float, optional): The fraction of available CUDA memory to use. Defaults to 0.60.
-        batch_size (int, optional): The default batch size to use if an error is detected. Defaults to 16.
+        model (torch.nn.Module): YOLO model to compute batch size for.
+        imgsz (int, optional): The image size used as input for the YOLO model.
+        fraction (float, optional): The fraction of available CUDA memory to use.
+        batch_size (int, optional): The default batch size to use if an error is detected.
         max_num_obj (int, optional): The maximum number of objects from dataset.
 
     Returns: