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

ultralytics/engine/predictor.py

@@ -65,28 +65,54 @@ Example:
 
 class BasePredictor:
     """
-    BasePredictor.
-
     A base class for creating predictors.
 
+    This class provides the foundation for prediction functionality, handling model setup, inference,
+    and result processing across various input sources.
+
     Attributes:
         args (SimpleNamespace): Configuration for the predictor.
         save_dir (Path): Directory to save results.
         done_warmup (bool): Whether the predictor has finished setup.
-        model (nn.Module): Model used for prediction.
+        model (torch.nn.Module): Model used for prediction.
         data (dict): Data configuration.
         device (torch.device): Device used for prediction.
         dataset (Dataset): Dataset used for prediction.
-        vid_writer (dict): Dictionary of {save_path: video_writer, ...} writer for saving video output.
+        vid_writer (dict): Dictionary of {save_path: video_writer} for saving video output.
+        plotted_img (numpy.ndarray): Last plotted image.
+        source_type (SimpleNamespace): Type of input source.
+        seen (int): Number of images processed.
+        windows (List): List of window names for visualization.
+        batch (tuple): Current batch data.
+        results (List): Current batch results.
+        transforms (callable): Image transforms for classification.
+        callbacks (dict): Callback functions for different events.
+        txt_path (Path): Path to save text results.
+        _lock (threading.Lock): Lock for thread-safe inference.
+
+    Methods:
+        preprocess: Prepare input image before inference.
+        inference: Run inference on a given image.
+        postprocess: Process raw predictions into structured results.
+        predict_cli: Run prediction for command line interface.
+        setup_source: Set up input source and inference mode.
+        stream_inference: Stream inference on input source.
+        setup_model: Initialize and configure the model.
+        write_results: Write inference results to files.
+        save_predicted_images: Save prediction visualizations.
+        show: Display results in a window.
+        run_callbacks: Execute registered callbacks for an event.
+        add_callback: Register a new callback function.
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
         """
-        Initializes the BasePredictor class.
+        Initialize the BasePredictor class.
 
         Args:
-            cfg (str, optional): Path to a configuration file. Defaults to DEFAULT_CFG.
-            overrides (dict, optional): Configuration overrides. Defaults to None.
+            cfg (str | dict): Path to a configuration file or a configuration dictionary.
+            overrides (dict | None): Configuration overrides.
+            _callbacks (dict | None): Dictionary of callback functions.
         """
         self.args = get_cfg(cfg, overrides)
         self.save_dir = get_save_dir(self.args)
@@ -120,7 +146,7 @@ class BasePredictor:
         Prepares input image before inference.
 
         Args:
-            im (torch.Tensor | List(np.ndarray)): BCHW for tensor, [(HWC) x B] for list.
+            im (torch.Tensor | List(np.ndarray)): Images of shape (N, 3, h, w) for tensor, [(h, w, 3) x N] for list.
         """
         not_tensor = not isinstance(im, torch.Tensor)
         if not_tensor:
@@ -136,7 +162,7 @@ class BasePredictor:
         return im
 
     def inference(self, im, *args, **kwargs):
-        """Runs inference on a given image using the specified model and arguments."""
+        """Run inference on a given image using the specified model and arguments."""
         visualize = (
             increment_path(self.save_dir / Path(self.batch[0][0]).stem, mkdir=True)
             if self.args.visualize and (not self.source_type.tensor)
@@ -149,10 +175,10 @@ class BasePredictor:
         Pre-transform input image before inference.
 
         Args:
-            im (List(np.ndarray)): (N, 3, h, w) for tensor, [(h, w, 3) x N] for list.
+            im (List[np.ndarray]): Images of shape (N, 3, h, w) for tensor, [(h, w, 3) x N] for list.
 
         Returns:
-            (list): A list of transformed images.
+            (List[np.ndarray]): A list of transformed images.
         """
         same_shapes = len({x.shape for x in im}) == 1
         letterbox = LetterBox(
@@ -163,11 +189,24 @@ class BasePredictor:
         return [letterbox(image=x) for x in im]
 
     def postprocess(self, preds, img, orig_imgs):
-        """Post-processes predictions for an image and returns them."""
+        """Post-process predictions for an image and return them."""
         return preds
 
     def __call__(self, source=None, model=None, stream=False, *args, **kwargs):
-        """Performs inference on an image or stream."""
+        """
+        Perform inference on an image or stream.
+
+        Args:
+            source (str | Path | List[str] | List[Path] | List[np.ndarray] | np.ndarray | torch.Tensor | None):
+                Source for inference.
+            model (str | Path | torch.nn.Module | None): Model for inference.
+            stream (bool): Whether to stream the inference results. If True, returns a generator.
+            *args (Any): Additional arguments for the inference method.
+            **kwargs (Any): Additional keyword arguments for the inference method.
+
+        Returns:
+            (List[ultralytics.engine.results.Results] | generator): Results objects or generator of Results objects.
+        """
         self.stream = stream
         if stream:
             return self.stream_inference(source, model, *args, **kwargs)
@@ -182,6 +221,11 @@ class BasePredictor:
         the inputs in a streaming manner. This method ensures that no outputs accumulate in memory by consuming the
         generator without storing results.
 
+        Args:
+            source (str | Path | List[str] | List[Path] | List[np.ndarray] | np.ndarray | torch.Tensor | None):
+                Source for inference.
+            model (str | Path | torch.nn.Module | None): Model for inference.
+
         Note:
             Do not modify this function or remove the generator. The generator ensures that no outputs are
             accumulated in memory, which is critical for preventing memory issues during long-running predictions.
@@ -191,7 +235,13 @@ class BasePredictor:
             pass
 
     def setup_source(self, source):
-        """Sets up source and inference mode."""
+        """
+        Set up source and inference mode.
+
+        Args:
+            source (str | Path | List[str] | List[Path] | List[np.ndarray] | np.ndarray | torch.Tensor):
+                Source for inference.
+        """
         self.imgsz = check_imgsz(self.args.imgsz, stride=self.model.stride, min_dim=2)  # check image size
         self.transforms = (
             getattr(
@@ -220,7 +270,19 @@ class BasePredictor:
 
     @smart_inference_mode()
     def stream_inference(self, source=None, model=None, *args, **kwargs):
-        """Streams real-time inference on camera feed and saves results to file."""
+        """
+        Stream real-time inference on camera feed and save results to file.
+
+        Args:
+            source (str | Path | List[str] | List[Path] | List[np.ndarray] | np.ndarray | torch.Tensor | None):
+                Source for inference.
+            model (str | Path | torch.nn.Module | None): Model for inference.
+            *args (Any): Additional arguments for the inference method.
+            **kwargs (Any): Additional keyword arguments for the inference method.
+
+        Yields:
+            (ultralytics.engine.results.Results): Results objects.
+        """
         if self.args.verbose:
             LOGGER.info("")
 
@@ -306,7 +368,13 @@ class BasePredictor:
         self.run_callbacks("on_predict_end")
 
     def setup_model(self, model, verbose=True):
-        """Initialize YOLO model with given parameters and set it to evaluation mode."""
+        """
+        Initialize YOLO model with given parameters and set it to evaluation mode.
+
+        Args:
+            model (str | Path | torch.nn.Module | None): Model to load or use.
+            verbose (bool): Whether to print verbose output.
+        """
         self.model = AutoBackend(
             weights=model or self.args.model,
             device=select_device(self.args.device, verbose=verbose),
@@ -323,7 +391,18 @@ class BasePredictor:
         self.model.eval()
 
     def write_results(self, i, p, im, s):
-        """Write inference results to a file or directory."""
+        """
+        Write inference results to a file or directory.
+
+        Args:
+            i (int): Index of the current image in the batch.
+            p (Path): Path to the current image.
+            im (torch.Tensor): Preprocessed image tensor.
+            s (List[str]): List of result strings.
+
+        Returns:
+            (str): String with result information.
+        """
         string = ""  # print string
         if len(im.shape) == 3:
             im = im[None]  # expand for batch dim
@@ -363,7 +442,13 @@ class BasePredictor:
         return string
 
     def save_predicted_images(self, save_path="", frame=0):
-        """Save video predictions as mp4 at specified path."""
+        """
+        Save video predictions as mp4 or images as jpg at specified path.
+
+        Args:
+            save_path (str): Path to save the results.
+            frame (int): Frame number for video mode.
+        """
         im = self.plotted_img
 
         # Save videos and streams
@@ -391,7 +476,7 @@ class BasePredictor:
             cv2.imwrite(str(Path(save_path).with_suffix(".jpg")), im)  # save to JPG for best support
 
     def show(self, p=""):
-        """Display an image in a window using the OpenCV imshow function."""
+        """Display an image in a window."""
         im = self.plotted_img
         if platform.system() == "Linux" and p not in self.windows:
             self.windows.append(p)
@@ -401,10 +486,10 @@ class BasePredictor:
         cv2.waitKey(300 if self.dataset.mode == "image" else 1)  # 1 millisecond
 
     def run_callbacks(self, event: str):
-        """Runs all registered callbacks for a specific event."""
+        """Run all registered callbacks for a specific event."""
         for callback in self.callbacks.get(event, []):
             callback(self)
 
     def add_callback(self, event: str, func):
-        """Add callback."""
+        """Add a callback function for a specific event."""
         self.callbacks[event].append(func)

ultralytics/engine/trainer.py

@@ -87,17 +87,20 @@ class BaseTrainer:
         fitness (float): Current fitness value.
         loss (float): Current loss value.
         tloss (float): Total loss value.
-        loss_names (list): List of loss names.
+        loss_names (List): List of loss names.
         csv (Path): Path to results CSV file.
+        metrics (Dict): Dictionary of metrics.
+        plots (Dict): Dictionary of plots.
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
         """
-        Initializes the BaseTrainer class.
+        Initialize the BaseTrainer class.
 
         Args:
             cfg (str, optional): Path to a configuration file. Defaults to DEFAULT_CFG.
-            overrides (dict, optional): Configuration overrides. Defaults to None.
+            overrides (Dict, optional): Configuration overrides. Defaults to None.
+            _callbacks (List, optional): List of callback functions. Defaults to None.
         """
         self.args = get_cfg(cfg, overrides)
         self.check_resume(overrides)
@@ -156,11 +159,11 @@ class BaseTrainer:
             callbacks.add_integration_callbacks(self)
 
     def add_callback(self, event: str, callback):
-        """Appends the given callback."""
+        """Append the given callback to the event's callback list."""
         self.callbacks[event].append(callback)
 
     def set_callback(self, event: str, callback):
-        """Overrides the existing callbacks with the given callback."""
+        """Override the existing callbacks with the given callback for the specified event."""
         self.callbacks[event] = [callback]
 
     def run_callbacks(self, event: str):
@@ -216,7 +219,7 @@ class BaseTrainer:
         self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=self.lf)
 
     def _setup_ddp(self, world_size):
-        """Initializes and sets the DistributedDataParallel parameters for training."""
+        """Initialize and set the DistributedDataParallel parameters for training."""
         torch.cuda.set_device(RANK)
         self.device = torch.device("cuda", RANK)
         # LOGGER.info(f'DDP info: RANK {RANK}, WORLD_SIZE {world_size}, DEVICE {self.device}')
@@ -229,7 +232,7 @@ class BaseTrainer:
         )
 
     def _setup_train(self, world_size):
-        """Builds dataloaders and optimizer on correct rank process."""
+        """Build dataloaders and optimizer on correct rank process."""
         # Model
         self.run_callbacks("on_pretrain_routine_start")
         ckpt = self.setup_model()
@@ -317,7 +320,7 @@ class BaseTrainer:
         self.run_callbacks("on_pretrain_routine_end")
 
     def _do_train(self, world_size=1):
-        """Train completed, evaluate and plot if specified by arguments."""
+        """Train the model with the specified world size."""
         if world_size > 1:
             self._setup_ddp(world_size)
         self._setup_train(world_size)
@@ -477,7 +480,7 @@ class BaseTrainer:
         self.run_callbacks("teardown")
 
     def auto_batch(self, max_num_obj=0):
-        """Get batch size by calculating memory occupation of model."""
+        """Calculate optimal batch size based on model and device memory constraints."""
         return check_train_batch_size(
             model=self.model,
             imgsz=self.args.imgsz,
@@ -487,12 +490,12 @@ class BaseTrainer:
         )  # returns batch size
 
     def _get_memory(self, fraction=False):
-        """Get accelerator memory utilization in GB or fraction."""
+        """Get accelerator memory utilization in GB or as a fraction of total memory."""
         memory, total = 0, 0
         if self.device.type == "mps":
             memory = torch.mps.driver_allocated_memory()
             if fraction:
-                total = torch.mps.get_mem_info()[0]
+                return __import__("psutil").virtual_memory().percent / 100
         elif self.device.type == "cpu":
             pass
         else:
@@ -502,7 +505,7 @@ class BaseTrainer:
         return ((memory / total) if total > 0 else 0) if fraction else (memory / 2**30)
 
     def _clear_memory(self):
-        """Clear accelerator memory on different platforms."""
+        """Clear accelerator memory by calling garbage collector and emptying cache."""
         gc.collect()
         if self.device.type == "mps":
             torch.mps.empty_cache()
@@ -512,7 +515,7 @@ class BaseTrainer:
             torch.cuda.empty_cache()
 
     def read_results_csv(self):
-        """Read results.csv into a dict using pandas."""
+        """Read results.csv into a dictionary using pandas."""
         import pandas as pd  # scope for faster 'import ultralytics'
 
         return pd.read_csv(self.csv).to_dict(orient="list")
@@ -554,9 +557,10 @@ class BaseTrainer:
 
     def get_dataset(self):
         """
-        Get train, val path from data dict if it exists.
+        Get train and validation datasets from data dictionary.
 
-        Returns None if data format is not recognized.
+        Returns:
+            (tuple): A tuple containing the training and validation/test datasets.
         """
         try:
             if self.args.task == "classify":
@@ -580,7 +584,12 @@ class BaseTrainer:
         return data["train"], data.get("val") or data.get("test")
 
     def setup_model(self):
-        """Load/create/download model for any task."""
+        """
+        Load, create, or download model for any task.
+
+        Returns:
+            (dict): Optional checkpoint to resume training from.
+        """
         if isinstance(self.model, torch.nn.Module):  # if model is loaded beforehand. No setup needed
             return
 
@@ -610,9 +619,10 @@ class BaseTrainer:
 
     def validate(self):
         """
-        Runs validation on test set using self.validator.
+        Run validation on test set using self.validator.
 
-        The returned dict is expected to contain "fitness" key.
+        Returns:
+            (tuple): A tuple containing metrics dictionary and fitness score.
         """
         metrics = self.validator(self)
         fitness = metrics.pop("fitness", -self.loss.detach().cpu().numpy())  # use loss as fitness measure if not found
@@ -646,7 +656,7 @@ class BaseTrainer:
         return {"loss": loss_items} if loss_items is not None else ["loss"]
 
     def set_model_attributes(self):
-        """To set or update model parameters before training."""
+        """Set or update model parameters before training."""
         self.model.names = self.data["names"]
 
     def build_targets(self, preds, targets):
@@ -667,7 +677,7 @@ class BaseTrainer:
         pass
 
     def save_metrics(self, metrics):
-        """Saves training metrics to a CSV file."""
+        """Save training metrics to a CSV file."""
         keys, vals = list(metrics.keys()), list(metrics.values())
         n = len(metrics) + 2  # number of cols
         s = "" if self.csv.exists() else (("%s," * n % tuple(["epoch", "time"] + keys)).rstrip(",") + "\n")  # header
@@ -685,7 +695,7 @@ class BaseTrainer:
         self.plots[path] = {"data": data, "timestamp": time.time()}
 
     def final_eval(self):
-        """Performs final evaluation and validation for object detection YOLO model."""
+        """Perform final evaluation and validation for object detection YOLO model."""
         ckpt = {}
         for f in self.last, self.best:
             if f.exists():
@@ -769,8 +779,7 @@ class BaseTrainer:
 
     def build_optimizer(self, model, name="auto", lr=0.001, momentum=0.9, decay=1e-5, iterations=1e5):
         """
-        Constructs an optimizer for the given model, based on the specified optimizer name, learning rate, momentum,
-        weight decay, and number of iterations.
+        Construct an optimizer for the given model.
 
         Args:
             model (torch.nn.Module): The model for which to build an optimizer.

ultralytics/engine/tuner.py

@@ -7,14 +7,11 @@ Hyperparameter tuning is the process of systematically searching for the optimal
 that yield the best model performance. This is particularly crucial in deep learning models like YOLO,
 where small changes in hyperparameters can lead to significant differences in model accuracy and efficiency.
 
-Example:
+Examples:
     Tune hyperparameters for YOLO11n on COCO8 at imgsz=640 and epochs=30 for 300 tuning iterations.
-    ```python
-    from ultralytics import YOLO
-
-    model = YOLO("yolo11n.pt")
-    model.tune(data="coco8.yaml", epochs=10, iterations=300, optimizer="AdamW", plots=False, save=False, val=False)
-    ```
+    >>> from ultralytics import YOLO
+    >>> model = YOLO("yolo11n.pt")
+    >>> model.tune(data="coco8.yaml", epochs=10, iterations=300, optimizer="AdamW", plots=False, save=False, val=False)
 """
 
 import random
@@ -32,39 +29,33 @@ from ultralytics.utils.plotting import plot_tune_results
 
 class Tuner:
     """
-    Class responsible for hyperparameter tuning of YOLO models.
+    A class for hyperparameter tuning of YOLO models.
 
-    The class evolves YOLO model hyperparameters over a given number of iterations
-    by mutating them according to the search space and retraining the model to evaluate their performance.
+    The class evolves YOLO model hyperparameters over a given number of iterations by mutating them according to the
+    search space and retraining the model to evaluate their performance.
 
     Attributes:
-        space (dict): Hyperparameter search space containing bounds and scaling factors for mutation.
+        space (Dict): Hyperparameter search space containing bounds and scaling factors for mutation.
         tune_dir (Path): Directory where evolution logs and results will be saved.
         tune_csv (Path): Path to the CSV file where evolution logs are saved.
+        args (Dict): Configuration arguments for the tuning process.
+        callbacks (List): Callback functions to be executed during tuning.
+        prefix (str): Prefix string for logging messages.
 
     Methods:
-        _mutate(hyp: dict) -> dict:
-            Mutates the given hyperparameters within the bounds specified in `self.space`.
-
-        __call__():
-            Executes the hyperparameter evolution across multiple iterations.
+        _mutate: Mutates the given hyperparameters within the specified bounds.
+        __call__: Executes the hyperparameter evolution across multiple iterations.
 
-    Example:
+    Examples:
         Tune hyperparameters for YOLO11n on COCO8 at imgsz=640 and epochs=30 for 300 tuning iterations.
-        ```python
-        from ultralytics import YOLO
-
-        model = YOLO("yolo11n.pt")
-        model.tune(data="coco8.yaml", epochs=10, iterations=300, optimizer="AdamW", plots=False, save=False, val=False)
-        ```
+        >>> from ultralytics import YOLO
+        >>> model = YOLO("yolo11n.pt")
+        >>> model.tune(
+        ...     data="coco8.yaml", epochs=10, iterations=300, optimizer="AdamW", plots=False, save=False, val=False
+        ... )
 
         Tune with custom search space.
-        ```python
-        from ultralytics import YOLO
-
-        model = YOLO("yolo11n.pt")
-        model.tune(space={key1: val1, key2: val2})  # custom search space dictionary
-        ```
+        >>> model.tune(space={key1: val1, key2: val2})  # custom search space dictionary
     """
 
     def __init__(self, args=DEFAULT_CFG, _callbacks=None):
@@ -72,7 +63,8 @@ class Tuner:
         Initialize the Tuner with configurations.
 
         Args:
-            args (dict, optional): Configuration for hyperparameter evolution.
+            args (Dict): Configuration for hyperparameter evolution.
+            _callbacks (List, optional): Callback functions to be executed during tuning.
         """
         self.space = args.pop("space", None) or {  # key: (min, max, gain(optional))
             # 'optimizer': tune.choice(['SGD', 'Adam', 'AdamW', 'NAdam', 'RAdam', 'RMSProp']),
@@ -114,7 +106,7 @@ class Tuner:
 
     def _mutate(self, parent="single", n=5, mutation=0.8, sigma=0.2):
         """
-        Mutates the hyperparameters based on bounds and scaling factors specified in `self.space`.
+        Mutate hyperparameters based on bounds and scaling factors specified in `self.space`.
 
         Args:
             parent (str): Parent selection method: 'single' or 'weighted'.
@@ -123,7 +115,7 @@ class Tuner:
             sigma (float): Standard deviation for Gaussian random number generator.
 
         Returns:
-            (dict): A dictionary containing mutated hyperparameters.
+            (Dict): A dictionary containing mutated hyperparameters.
         """
         if self.tune_csv.exists():  # if CSV file exists: select best hyps and mutate
             # Select parent(s)
@@ -160,22 +152,23 @@ class Tuner:
 
     def __call__(self, model=None, iterations=10, cleanup=True):
         """
-        Executes the hyperparameter evolution process when the Tuner instance is called.
+        Execute the hyperparameter evolution process when the Tuner instance is called.
 
         This method iterates through the number of iterations, performing the following steps in each iteration:
+
         1. Load the existing hyperparameters or initialize new ones.
         2. Mutate the hyperparameters using the `mutate` method.
         3. Train a YOLO model with the mutated hyperparameters.
         4. Log the fitness score and mutated hyperparameters to a CSV file.
 
         Args:
-           model (Model): A pre-initialized YOLO model to be used for training.
-           iterations (int): The number of generations to run the evolution for.
-           cleanup (bool): Whether to delete iteration weights to reduce storage space used during tuning.
+            model (Model): A pre-initialized YOLO model to be used for training.
+            iterations (int): The number of generations to run the evolution for.
+            cleanup (bool): Whether to delete iteration weights to reduce storage space used during tuning.
 
         Note:
-           The method utilizes the `self.tune_csv` Path object to read and log hyperparameters and fitness scores.
-           Ensure this path is set correctly in the Tuner instance.
+            The method utilizes the `self.tune_csv` Path object to read and log hyperparameters and fitness scores.
+            Ensure this path is set correctly in the Tuner instance.
         """
         t0 = time.time()
         best_save_dir, best_metrics = None, None