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

ultralytics/models/yolo/classify/val.py

@@ -13,21 +13,43 @@ class ClassificationValidator(BaseValidator):
     """
     A class extending the BaseValidator class for validation based on a classification model.
 
-    Notes:
-        - Torchvision classification models can also be passed to the 'model' argument, i.e. model='resnet18'.
-
-    Example:
-        ```python
-        from ultralytics.models.yolo.classify import ClassificationValidator
+    This validator handles the validation process for classification models, including metrics calculation,
+    confusion matrix generation, and visualization of results.
+
+    Attributes:
+        targets (List[torch.Tensor]): Ground truth class labels.
+        pred (List[torch.Tensor]): Model predictions.
+        metrics (ClassifyMetrics): Object to calculate and store classification metrics.
+        names (Dict): Mapping of class indices to class names.
+        nc (int): Number of classes.
+        confusion_matrix (ConfusionMatrix): Matrix to evaluate model performance across classes.
+
+    Methods:
+        get_desc: Return a formatted string summarizing classification metrics.
+        init_metrics: Initialize confusion matrix, class names, and tracking containers.
+        preprocess: Preprocess input batch by moving data to device.
+        update_metrics: Update running metrics with model predictions and batch targets.
+        finalize_metrics: Finalize metrics including confusion matrix and processing speed.
+        postprocess: Extract the primary prediction from model output.
+        get_stats: Calculate and return a dictionary of metrics.
+        build_dataset: Create a ClassificationDataset instance for validation.
+        get_dataloader: Build and return a data loader for classification validation.
+        print_results: Print evaluation metrics for the classification model.
+        plot_val_samples: Plot validation image samples with their ground truth labels.
+        plot_predictions: Plot images with their predicted class labels.
+
+    Examples:
+        >>> from ultralytics.models.yolo.classify import ClassificationValidator
+        >>> args = dict(model="yolo11n-cls.pt", data="imagenet10")
+        >>> validator = ClassificationValidator(args=args)
+        >>> validator()
 
-        args = dict(model="yolo11n-cls.pt", data="imagenet10")
-        validator = ClassificationValidator(args=args)
-        validator()
-        ```
+    Notes:
+        Torchvision classification models can also be passed to the 'model' argument, i.e. model='resnet18'.
     """
 
     def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None):
-        """Initializes ClassificationValidator instance with args, dataloader, save_dir, and progress bar."""
+        """Initialize ClassificationValidator with dataloader, save directory, and other parameters."""
         super().__init__(dataloader, save_dir, pbar, args, _callbacks)
         self.targets = None
         self.pred = None
@@ -35,11 +57,11 @@ class ClassificationValidator(BaseValidator):
         self.metrics = ClassifyMetrics()
 
     def get_desc(self):
-        """Returns a formatted string summarizing classification metrics."""
+        """Return a formatted string summarizing classification metrics."""
         return ("%22s" + "%11s" * 2) % ("classes", "top1_acc", "top5_acc")
 
     def init_metrics(self, model):
-        """Initialize confusion matrix, class names, and top-1 and top-5 accuracy."""
+        """Initialize confusion matrix, class names, and tracking containers for predictions and targets."""
         self.names = model.names
         self.nc = len(model.names)
         self.confusion_matrix = ConfusionMatrix(nc=self.nc, conf=self.args.conf, task="classify")
@@ -47,20 +69,20 @@ class ClassificationValidator(BaseValidator):
         self.targets = []
 
     def preprocess(self, batch):
-        """Preprocesses input batch and returns it."""
+        """Preprocess input batch by moving data to device and converting to appropriate dtype."""
         batch["img"] = batch["img"].to(self.device, non_blocking=True)
         batch["img"] = batch["img"].half() if self.args.half else batch["img"].float()
         batch["cls"] = batch["cls"].to(self.device)
         return batch
 
     def update_metrics(self, preds, batch):
-        """Updates running metrics with model predictions and batch targets."""
+        """Update running metrics with model predictions and batch targets."""
         n5 = min(len(self.names), 5)
         self.pred.append(preds.argsort(1, descending=True)[:, :n5].type(torch.int32).cpu())
         self.targets.append(batch["cls"].type(torch.int32).cpu())
 
     def finalize_metrics(self, *args, **kwargs):
-        """Finalizes metrics of the model such as confusion_matrix and speed."""
+        """Finalize metrics including confusion matrix and processing speed."""
         self.confusion_matrix.process_cls_preds(self.pred, self.targets)
         if self.args.plots:
             for normalize in True, False:
@@ -72,30 +94,30 @@ class ClassificationValidator(BaseValidator):
         self.metrics.save_dir = self.save_dir
 
     def postprocess(self, preds):
-        """Preprocesses the classification predictions."""
+        """Extract the primary prediction from model output if it's in a list or tuple format."""
         return preds[0] if isinstance(preds, (list, tuple)) else preds
 
     def get_stats(self):
-        """Returns a dictionary of metrics obtained by processing targets and predictions."""
+        """Calculate and return a dictionary of metrics by processing targets and predictions."""
         self.metrics.process(self.targets, self.pred)
         return self.metrics.results_dict
 
     def build_dataset(self, img_path):
-        """Creates and returns a ClassificationDataset instance using given image path and preprocessing parameters."""
+        """Create a ClassificationDataset instance for validation."""
         return ClassificationDataset(root=img_path, args=self.args, augment=False, prefix=self.args.split)
 
     def get_dataloader(self, dataset_path, batch_size):
-        """Builds and returns a data loader for classification tasks with given parameters."""
+        """Build and return a data loader for classification validation."""
         dataset = self.build_dataset(dataset_path)
         return build_dataloader(dataset, batch_size, self.args.workers, rank=-1)
 
     def print_results(self):
-        """Prints evaluation metrics for YOLO object detection model."""
+        """Print evaluation metrics for the classification model."""
         pf = "%22s" + "%11.3g" * len(self.metrics.keys)  # print format
         LOGGER.info(pf % ("all", self.metrics.top1, self.metrics.top5))
 
     def plot_val_samples(self, batch, ni):
-        """Plot validation image samples."""
+        """Plot validation image samples with their ground truth labels."""
         plot_images(
             images=batch["img"],
             batch_idx=torch.arange(len(batch["img"])),
@@ -106,7 +128,7 @@ class ClassificationValidator(BaseValidator):
         )
 
     def plot_predictions(self, batch, preds, ni):
-        """Plots predicted bounding boxes on input images and saves the result."""
+        """Plot images with their predicted class labels and save the visualization."""
         plot_images(
             batch["img"],
             batch_idx=torch.arange(len(batch["img"])),

ultralytics/models/yolo/detect/predict.py

@@ -9,15 +9,25 @@ class DetectionPredictor(BasePredictor):
     """
     A class extending the BasePredictor class for prediction based on a detection model.
 
-    Example:
-        ```python
-        from ultralytics.utils import ASSETS
-        from ultralytics.models.yolo.detect import DetectionPredictor
-
-        args = dict(model="yolo11n.pt", source=ASSETS)
-        predictor = DetectionPredictor(overrides=args)
-        predictor.predict_cli()
-        ```
+    This predictor specializes in object detection tasks, processing model outputs into meaningful detection results
+    with bounding boxes and class predictions.
+
+    Attributes:
+        args (namespace): Configuration arguments for the predictor.
+        model (nn.Module): The detection model used for inference.
+        batch (List): Batch of images and metadata for processing.
+
+    Methods:
+        postprocess: Process raw model predictions into detection results.
+        construct_results: Build Results objects from processed predictions.
+        construct_result: Create a single Result object from a prediction.
+
+    Examples:
+        >>> from ultralytics.utils import ASSETS
+        >>> from ultralytics.models.yolo.detect import DetectionPredictor
+        >>> args = dict(model="yolo11n.pt", source=ASSETS)
+        >>> predictor = DetectionPredictor(overrides=args)
+        >>> predictor.predict_cli()
     """
 
     def postprocess(self, preds, img, orig_imgs, **kwargs):
@@ -41,15 +51,15 @@ class DetectionPredictor(BasePredictor):
 
     def construct_results(self, preds, img, orig_imgs):
         """
-        Constructs a list of result objects from the predictions.
+        Construct a list of Results objects from model predictions.
 
         Args:
-            preds (List[torch.Tensor]): List of predicted bounding boxes and scores.
-            img (torch.Tensor): The image after preprocessing.
+            preds (List[torch.Tensor]): List of predicted bounding boxes and scores for each image.
+            img (torch.Tensor): Batch of preprocessed images used for inference.
             orig_imgs (List[np.ndarray]): List of original images before preprocessing.
 
         Returns:
-            (list): List of result objects containing the original images, image paths, class names, and bounding boxes.
+            (List[Results]): List of Results objects containing detection information for each image.
         """
         return [
             self.construct_result(pred, img, orig_img, img_path)
@@ -58,16 +68,16 @@ class DetectionPredictor(BasePredictor):
 
     def construct_result(self, pred, img, orig_img, img_path):
         """
-        Constructs the result object from the prediction.
+        Construct a single Results object from one image prediction.
 
         Args:
-            pred (torch.Tensor): The predicted bounding boxes and scores.
-            img (torch.Tensor): The image after preprocessing.
-            orig_img (np.ndarray): The original image before preprocessing.
-            img_path (str): The path to the original image.
+            pred (torch.Tensor): Predicted boxes and scores with shape (N, 6) where N is the number of detections.
+            img (torch.Tensor): Preprocessed image tensor used for inference.
+            orig_img (np.ndarray): Original image before preprocessing.
+            img_path (str): Path to the original image file.
 
         Returns:
-            (Results): The result object containing the original image, image path, class names, and bounding boxes.
+            (Results): Results object containing the original image, image path, class names, and scaled bounding boxes.
         """
         pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
         return Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6])

ultralytics/models/yolo/detect/train.py

@@ -20,30 +20,63 @@ class DetectionTrainer(BaseTrainer):
     """
     A class extending the BaseTrainer class for training based on a detection model.
 
-    Example:
-        ```python
-        from ultralytics.models.yolo.detect import DetectionTrainer
-
-        args = dict(model="yolo11n.pt", data="coco8.yaml", epochs=3)
-        trainer = DetectionTrainer(overrides=args)
-        trainer.train()
-        ```
+    This trainer specializes in object detection tasks, handling the specific requirements for training YOLO models
+    for object detection.
+
+    Attributes:
+        model (DetectionModel): The YOLO detection model being trained.
+        data (Dict): Dictionary containing dataset information including class names and number of classes.
+        loss_names (Tuple[str]): Names of the loss components used in training (box_loss, cls_loss, dfl_loss).
+
+    Methods:
+        build_dataset: Build YOLO dataset for training or validation.
+        get_dataloader: Construct and return dataloader for the specified mode.
+        preprocess_batch: Preprocess a batch of images by scaling and converting to float.
+        set_model_attributes: Set model attributes based on dataset information.
+        get_model: Return a YOLO detection model.
+        get_validator: Return a validator for model evaluation.
+        label_loss_items: Return a loss dictionary with labeled training loss items.
+        progress_string: Return a formatted string of training progress.
+        plot_training_samples: Plot training samples with their annotations.
+        plot_metrics: Plot metrics from a CSV file.
+        plot_training_labels: Create a labeled training plot of the YOLO model.
+        auto_batch: Calculate optimal batch size based on model memory requirements.
+
+    Examples:
+        >>> from ultralytics.models.yolo.detect import DetectionTrainer
+        >>> args = dict(model="yolo11n.pt", data="coco8.yaml", epochs=3)
+        >>> trainer = DetectionTrainer(overrides=args)
+        >>> trainer.train()
     """
 
     def build_dataset(self, img_path, mode="train", batch=None):
         """
-        Build YOLO Dataset.
+        Build YOLO Dataset for training or validation.
 
         Args:
             img_path (str): Path to the folder containing images.
             mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode.
-            batch (int, optional): Size of batches, this is for `rect`. Defaults to None.
+            batch (int, optional): Size of batches, this is for `rect`.
+
+        Returns:
+            (Dataset): YOLO dataset object configured for the specified mode.
         """
         gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
         return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs)
 
     def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
-        """Construct and return dataloader."""
+        """
+        Construct and return dataloader for the specified mode.
+
+        Args:
+            dataset_path (str): Path to the dataset.
+            batch_size (int): Number of images per batch.
+            rank (int): Process rank for distributed training.
+            mode (str): 'train' for training dataloader, 'val' for validation dataloader.
+
+        Returns:
+            (DataLoader): PyTorch dataloader object.
+        """
         assert mode in {"train", "val"}, f"Mode must be 'train' or 'val', not {mode}."
         with torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDP
             dataset = self.build_dataset(dataset_path, mode, batch_size)
@@ -55,7 +88,15 @@ class DetectionTrainer(BaseTrainer):
         return build_dataloader(dataset, batch_size, workers, shuffle, rank)  # return dataloader
 
     def preprocess_batch(self, batch):
-        """Preprocesses a batch of images by scaling and converting to float."""
+        """
+        Preprocess a batch of images by scaling and converting to float.
+
+        Args:
+            batch (Dict): Dictionary containing batch data with 'img' tensor.
+
+        Returns:
+            (Dict): Preprocessed batch with normalized images.
+        """
         batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
         if self.args.multi_scale:
             imgs = batch["img"]
@@ -74,7 +115,8 @@ class DetectionTrainer(BaseTrainer):
         return batch
 
     def set_model_attributes(self):
-        """Nl = de_parallel(self.model).model[-1].nl  # number of detection layers (to scale hyps)."""
+        """Set model attributes based on dataset information."""
+        # Nl = de_parallel(self.model).model[-1].nl  # number of detection layers (to scale hyps)
         # self.args.box *= 3 / nl  # scale to layers
         # self.args.cls *= self.data["nc"] / 80 * 3 / nl  # scale to classes and layers
         # self.args.cls *= (self.args.imgsz / 640) ** 2 * 3 / nl  # scale to image size and layers
@@ -84,14 +126,24 @@ class DetectionTrainer(BaseTrainer):
         # TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
 
     def get_model(self, cfg=None, weights=None, verbose=True):
-        """Return a YOLO detection model."""
+        """
+        Return a YOLO detection model.
+
+        Args:
+            cfg (str, optional): Path to model configuration file.
+            weights (str, optional): Path to model weights.
+            verbose (bool): Whether to display model information.
+
+        Returns:
+            (DetectionModel): YOLO detection model.
+        """
         model = DetectionModel(cfg, nc=self.data["nc"], verbose=verbose and RANK == -1)
         if weights:
             model.load(weights)
         return model
 
     def get_validator(self):
-        """Returns a DetectionValidator for YOLO model validation."""
+        """Return a DetectionValidator for YOLO model validation."""
         self.loss_names = "box_loss", "cls_loss", "dfl_loss"
         return yolo.detect.DetectionValidator(
             self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
@@ -99,9 +151,14 @@ class DetectionTrainer(BaseTrainer):
 
     def label_loss_items(self, loss_items=None, prefix="train"):
         """
-        Returns a loss dict with labelled training loss items tensor.
+        Return a loss dict with labeled training loss items tensor.
 
-        Not needed for classification but necessary for segmentation & detection
+        Args:
+            loss_items (List[float], optional): List of loss values.
+            prefix (str): Prefix for keys in the returned dictionary.
+
+        Returns:
+            (Dict | List): Dictionary of labeled loss items if loss_items is provided, otherwise list of keys.
         """
         keys = [f"{prefix}/{x}" for x in self.loss_names]
         if loss_items is not None:
@@ -111,7 +168,7 @@ class DetectionTrainer(BaseTrainer):
             return keys
 
     def progress_string(self):
-        """Returns a formatted string of training progress with epoch, GPU memory, loss, instances and size."""
+        """Return a formatted string of training progress with epoch, GPU memory, loss, instances and size."""
         return ("\n" + "%11s" * (4 + len(self.loss_names))) % (
             "Epoch",
             "GPU_mem",
@@ -121,7 +178,13 @@ class DetectionTrainer(BaseTrainer):
         )
 
     def plot_training_samples(self, batch, ni):
-        """Plots training samples with their annotations."""
+        """
+        Plot training samples with their annotations.
+
+        Args:
+            batch (Dict): Dictionary containing batch data.
+            ni (int): Number of iterations.
+        """
         plot_images(
             images=batch["img"],
             batch_idx=batch["batch_idx"],
@@ -133,7 +196,7 @@ class DetectionTrainer(BaseTrainer):
         )
 
     def plot_metrics(self):
-        """Plots metrics from a CSV file."""
+        """Plot metrics from a CSV file."""
         plot_results(file=self.csv, on_plot=self.on_plot)  # save results.png
 
     def plot_training_labels(self):
@@ -143,8 +206,12 @@ class DetectionTrainer(BaseTrainer):
         plot_labels(boxes, cls.squeeze(), names=self.data["names"], save_dir=self.save_dir, on_plot=self.on_plot)
 
     def auto_batch(self):
-        """Get batch size by calculating memory occupation of model."""
+        """
+        Get optimal batch size by calculating memory occupation of model.
+
+        Returns:
+            (int): Optimal batch size.
+        """
         train_dataset = self.build_dataset(self.trainset, mode="train", batch=16)
-        # 4 for mosaic augmentation
-        max_num_obj = max(len(label["cls"]) for label in train_dataset.labels) * 4
+        max_num_obj = max(len(label["cls"]) for label in train_dataset.labels) * 4  # 4 for mosaic augmentation
         return super().auto_batch(max_num_obj)