ultralytics 8.3.143__py3-none-any.whl → 8.3.145__py3-none-any.whl

ultralytics/models/rtdetr/train.py

@@ -1,6 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from copy import copy
+from typing import Optional
 
 from ultralytics.models.yolo.detect import DetectionTrainer
 from ultralytics.nn.tasks import RTDETRDetectionModel
@@ -18,12 +19,17 @@ class RTDETRTrainer(DetectionTrainer):
     speed.
 
     Attributes:
-        loss_names (Tuple[str]): Names of the loss components used for training.
+        loss_names (tuple): Names of the loss components used for training.
         data (dict): Dataset configuration containing class count and other parameters.
         args (dict): Training arguments and hyperparameters.
         save_dir (Path): Directory to save training results.
         test_loader (DataLoader): DataLoader for validation/testing data.
 
+    Methods:
+        get_model: Initialize and return an RT-DETR model for object detection tasks.
+        build_dataset: Build and return an RT-DETR dataset for training or validation.
+        get_validator: Return a DetectionValidator suitable for RT-DETR model validation.
+
     Notes:
         - F.grid_sample used in RT-DETR does not support the `deterministic=True` argument.
         - AMP training can lead to NaN outputs and may produce errors during bipartite graph matching.
@@ -35,7 +41,7 @@ class RTDETRTrainer(DetectionTrainer):
         >>> trainer.train()
     """
 
-    def get_model(self, cfg=None, weights=None, verbose=True):
+    def get_model(self, cfg: Optional[dict] = None, weights: Optional[str] = None, verbose: bool = True):
         """
         Initialize and return an RT-DETR model for object detection tasks.
 
@@ -52,7 +58,7 @@ class RTDETRTrainer(DetectionTrainer):
             model.load(weights)
         return model
 
-    def build_dataset(self, img_path, mode="val", batch=None):
+    def build_dataset(self, img_path: str, mode: str = "val", batch: Optional[int] = None):
         """
         Build and return an RT-DETR dataset for training or validation.
 
@@ -80,6 +86,6 @@ class RTDETRTrainer(DetectionTrainer):
         )
 
     def get_validator(self):
-        """Returns a DetectionValidator suitable for RT-DETR model validation."""
+        """Return a DetectionValidator suitable for RT-DETR model validation."""
         self.loss_names = "giou_loss", "cls_loss", "l1_loss"
         return RTDETRValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args))

ultralytics/models/rtdetr/val.py

@@ -16,6 +16,22 @@ class RTDETRDataset(YOLODataset):
 
     This specialized dataset class is designed for use with the RT-DETR object detection model and is optimized for
     real-time detection and tracking tasks.
+
+    Attributes:
+        augment (bool): Whether to apply data augmentation.
+        rect (bool): Whether to use rectangular training.
+        use_segments (bool): Whether to use segmentation masks.
+        use_keypoints (bool): Whether to use keypoint annotations.
+        imgsz (int): Target image size for training.
+
+    Methods:
+        load_image: Load one image from dataset index.
+        build_transforms: Build transformation pipeline for the dataset.
+
+    Examples:
+        Initialize an RT-DETR dataset
+        >>> dataset = RTDETRDataset(img_path="path/to/images", imgsz=640)
+        >>> image, hw = dataset.load_image(0)
     """
 
     def __init__(self, *args, data=None, **kwargs):
@@ -27,7 +43,7 @@ class RTDETRDataset(YOLODataset):
 
         Args:
             *args (Any): Variable length argument list passed to the parent YOLODataset class.
-            data (Dict | None): Dictionary containing dataset information. If None, default values will be used.
+            data (dict | None): Dictionary containing dataset information. If None, default values will be used.
             **kwargs (Any): Additional keyword arguments passed to the parent YOLODataset class.
         """
         super().__init__(*args, data=data, **kwargs)
@@ -41,11 +57,12 @@ class RTDETRDataset(YOLODataset):
             rect_mode (bool, optional): Whether to use rectangular mode for batch inference.
 
         Returns:
-            im (numpy.ndarray): The loaded image.
+            im (torch.Tensor): The loaded image.
             resized_hw (tuple): Height and width of the resized image with shape (2,).
 
         Examples:
-            >>> dataset = RTDETRDataset(...)
+            Load an image from the dataset
+            >>> dataset = RTDETRDataset(img_path="path/to/images")
             >>> image, hw = dataset.load_image(0)
         """
         return super().load_image(i=i, rect_mode=rect_mode)
@@ -90,13 +107,22 @@ class RTDETRValidator(DetectionValidator):
     The class allows building of an RTDETR-specific dataset for validation, applies Non-maximum suppression for
     post-processing, and updates evaluation metrics accordingly.
 
+    Attributes:
+        args (Namespace): Configuration arguments for validation.
+        data (dict): Dataset configuration dictionary.
+
+    Methods:
+        build_dataset: Build an RTDETR Dataset for validation.
+        postprocess: Apply Non-maximum suppression to prediction outputs.
+
     Examples:
+        Initialize and run RT-DETR validation
         >>> from ultralytics.models.rtdetr import RTDETRValidator
         >>> args = dict(model="rtdetr-l.pt", data="coco8.yaml")
         >>> validator = RTDETRValidator(args=args)
         >>> validator()
 
-    Note:
+    Notes:
         For further details on the attributes and methods, refer to the parent DetectionValidator class.
     """
 
@@ -106,7 +132,8 @@ class RTDETRValidator(DetectionValidator):
 
         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.
+            mode (str, optional): `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`.
 
         Returns:
@@ -129,10 +156,10 @@ class RTDETRValidator(DetectionValidator):
         Apply Non-maximum suppression to prediction outputs.
 
         Args:
-            preds (List | Tuple | torch.Tensor): Raw predictions from the model.
+            preds (list | tuple | torch.Tensor): Raw predictions from the model.
 
         Returns:
-            (List[torch.Tensor]): List of processed predictions for each image in batch.
+            (list[torch.Tensor]): List of processed predictions for each image in batch.
         """
         if not isinstance(preds, (list, tuple)):  # list for PyTorch inference but list[0] Tensor for export inference
             preds = [preds, None]
@@ -153,7 +180,7 @@ class RTDETRValidator(DetectionValidator):
 
     def _prepare_batch(self, si, batch):
         """
-        Prepares a batch for validation by applying necessary transformations.
+        Prepare a batch for validation by applying necessary transformations.
 
         Args:
             si (int): Batch index.
@@ -176,7 +203,7 @@ class RTDETRValidator(DetectionValidator):
 
     def _prepare_pred(self, pred, pbatch):
         """
-        Prepares predictions by scaling bounding boxes to original image dimensions.
+        Prepare predictions by scaling bounding boxes to original image dimensions.
 
         Args:
             pred (torch.Tensor): Raw predictions.

ultralytics/models/sam/amg.py

@@ -11,7 +11,24 @@ import torch
 def is_box_near_crop_edge(
     boxes: torch.Tensor, crop_box: List[int], orig_box: List[int], atol: float = 20.0
 ) -> torch.Tensor:
-    """Determines if bounding boxes are near the edge of a cropped image region using a specified tolerance."""
+    """
+    Determine if bounding boxes are near the edge of a cropped image region using a specified tolerance.
+
+    Args:
+        boxes (torch.Tensor): Bounding boxes in XYXY format.
+        crop_box (List[int]): Crop box coordinates in [x0, y0, x1, y1] format.
+        orig_box (List[int]): Original image box coordinates in [x0, y0, x1, y1] format.
+        atol (float, optional): Absolute tolerance for edge proximity detection.
+
+    Returns:
+        (torch.Tensor): Boolean tensor indicating which boxes are near crop edges.
+
+    Examples:
+        >>> boxes = torch.tensor([[10, 10, 50, 50], [100, 100, 150, 150]])
+        >>> crop_box = [0, 0, 200, 200]
+        >>> orig_box = [0, 0, 300, 300]
+        >>> near_edge = is_box_near_crop_edge(boxes, crop_box, orig_box, atol=20.0)
+    """
     crop_box_torch = torch.as_tensor(crop_box, dtype=torch.float, device=boxes.device)
     orig_box_torch = torch.as_tensor(orig_box, dtype=torch.float, device=boxes.device)
     boxes = uncrop_boxes_xyxy(boxes, crop_box).float()
@@ -52,7 +69,7 @@ def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]:
 
 def calculate_stability_score(masks: torch.Tensor, mask_threshold: float, threshold_offset: float) -> torch.Tensor:
     """
-    Computes the stability score for a batch of masks.
+    Compute the stability score for a batch of masks.
 
     The stability score is the IoU between binary masks obtained by thresholding the predicted mask logits at
     high and low values.
@@ -90,7 +107,7 @@ def build_point_grid(n_per_side: int) -> np.ndarray:
 
 
 def build_all_layer_point_grids(n_per_side: int, n_layers: int, scale_per_layer: int) -> List[np.ndarray]:
-    """Generates point grids for multiple crop layers with varying scales and densities."""
+    """Generate point grids for multiple crop layers with varying scales and densities."""
     return [build_point_grid(int(n_per_side / (scale_per_layer**i))) for i in range(n_layers + 1)]
 
 
@@ -98,7 +115,7 @@ def generate_crop_boxes(
     im_size: Tuple[int, ...], n_layers: int, overlap_ratio: float
 ) -> Tuple[List[List[int]], List[int]]:
     """
-    Generates crop boxes of varying sizes for multiscale image processing, with layered overlapping regions.
+    Generate crop boxes of varying sizes for multiscale image processing, with layered overlapping regions.
 
     Args:
         im_size (Tuple[int, ...]): Height and width of the input image.
@@ -106,8 +123,8 @@ def generate_crop_boxes(
         overlap_ratio (float): Ratio of overlap between adjacent crop boxes.
 
     Returns:
-        (List[List[int]]): List of crop boxes in [x0, y0, x1, y1] format.
-        (List[int]): List of layer indices corresponding to each crop box.
+        crop_boxes (List[List[int]]): List of crop boxes in [x0, y0, x1, y1] format.
+        layer_idxs (List[int]): List of layer indices corresponding to each crop box.
 
     Examples:
         >>> im_size = (800, 1200)  # Height, width
@@ -124,7 +141,7 @@ def generate_crop_boxes(
     layer_idxs.append(0)
 
     def crop_len(orig_len, n_crops, overlap):
-        """Calculates the length of each crop given the original length, number of crops, and overlap."""
+        """Calculate the length of each crop given the original length, number of crops, and overlap."""
         return int(math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops))
 
     for i_layer in range(n_layers):
@@ -179,16 +196,17 @@ def uncrop_masks(masks: torch.Tensor, crop_box: List[int], orig_h: int, orig_w:
 
 def remove_small_regions(mask: np.ndarray, area_thresh: float, mode: str) -> Tuple[np.ndarray, bool]:
     """
-    Removes small disconnected regions or holes in a mask based on area threshold and mode.
+    Remove small disconnected regions or holes in a mask based on area threshold and mode.
 
     Args:
         mask (np.ndarray): Binary mask to process.
         area_thresh (float): Area threshold below which regions will be removed.
-        mode (str): Processing mode, either 'holes' to fill small holes or 'islands' to remove small disconnected regions.
+        mode (str): Processing mode, either 'holes' to fill small holes or 'islands' to remove small disconnected
+            regions.
 
     Returns:
-        (np.ndarray): Processed binary mask with small regions removed.
-        (bool): Whether any regions were modified.
+        processed_mask (np.ndarray): Processed binary mask with small regions removed.
+        modified (bool): Whether any regions were modified.
 
     Examples:
         >>> mask = np.zeros((100, 100), dtype=np.bool_)
@@ -216,7 +234,7 @@ def remove_small_regions(mask: np.ndarray, area_thresh: float, mode: str) -> Tup
 
 def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor:
     """
-    Calculates bounding boxes in XYXY format around binary masks.
+    Calculate bounding boxes in XYXY format around binary masks.
 
     Args:
         masks (torch.Tensor): Binary masks with shape (B, H, W) or (B, C, H, W).

ultralytics/models/sam/build.py

@@ -21,7 +21,7 @@ from .modules.transformer import TwoWayTransformer
 
 
 def build_sam_vit_h(checkpoint=None):
-    """Builds and returns a Segment Anything Model (SAM) h-size model with specified encoder parameters."""
+    """Build and return a Segment Anything Model (SAM) h-size model with specified encoder parameters."""
     return _build_sam(
         encoder_embed_dim=1280,
         encoder_depth=32,
@@ -32,7 +32,7 @@ def build_sam_vit_h(checkpoint=None):
 
 
 def build_sam_vit_l(checkpoint=None):
-    """Builds and returns a Segment Anything Model (SAM) l-size model with specified encoder parameters."""
+    """Build and return a Segment Anything Model (SAM) l-size model with specified encoder parameters."""
     return _build_sam(
         encoder_embed_dim=1024,
         encoder_depth=24,
@@ -43,7 +43,7 @@ def build_sam_vit_l(checkpoint=None):
 
 
 def build_sam_vit_b(checkpoint=None):
-    """Constructs and returns a Segment Anything Model (SAM) with b-size architecture and optional checkpoint."""
+    """Build and return a Segment Anything Model (SAM) b-size model with specified encoder parameters."""
     return _build_sam(
         encoder_embed_dim=768,
         encoder_depth=12,
@@ -54,7 +54,7 @@ def build_sam_vit_b(checkpoint=None):
 
 
 def build_mobile_sam(checkpoint=None):
-    """Builds and returns a Mobile Segment Anything Model (Mobile-SAM) for efficient image segmentation."""
+    """Build and return a Mobile Segment Anything Model (Mobile-SAM) for efficient image segmentation."""
     return _build_sam(
         encoder_embed_dim=[64, 128, 160, 320],
         encoder_depth=[2, 2, 6, 2],
@@ -66,7 +66,7 @@ def build_mobile_sam(checkpoint=None):
 
 
 def build_sam2_t(checkpoint=None):
-    """Builds and returns a Segment Anything Model 2 (SAM2) tiny-size model with specified architecture parameters."""
+    """Build and return a Segment Anything Model 2 (SAM2) tiny-size model with specified architecture parameters."""
     return _build_sam2(
         encoder_embed_dim=96,
         encoder_stages=[1, 2, 7, 2],
@@ -79,7 +79,7 @@ def build_sam2_t(checkpoint=None):
 
 
 def build_sam2_s(checkpoint=None):
-    """Builds and returns a small-size Segment Anything Model (SAM2) with specified architecture parameters."""
+    """Build and return a small-size Segment Anything Model 2 (SAM2) with specified architecture parameters."""
     return _build_sam2(
         encoder_embed_dim=96,
         encoder_stages=[1, 2, 11, 2],
@@ -92,7 +92,7 @@ def build_sam2_s(checkpoint=None):
 
 
 def build_sam2_b(checkpoint=None):
-    """Builds and returns a SAM2 base-size model with specified architecture parameters."""
+    """Build and return a Segment Anything Model 2 (SAM2) base-size model with specified architecture parameters."""
     return _build_sam2(
         encoder_embed_dim=112,
         encoder_stages=[2, 3, 16, 3],
@@ -106,7 +106,7 @@ def build_sam2_b(checkpoint=None):
 
 
 def build_sam2_l(checkpoint=None):
-    """Builds and returns a large-size Segment Anything Model (SAM2) with specified architecture parameters."""
+    """Build and return a large-size Segment Anything Model 2 (SAM2) with specified architecture parameters."""
     return _build_sam2(
         encoder_embed_dim=144,
         encoder_stages=[2, 6, 36, 4],
@@ -127,15 +127,15 @@ def _build_sam(
     mobile_sam=False,
 ):
     """
-    Builds a Segment Anything Model (SAM) with specified encoder parameters.
+    Build a Segment Anything Model (SAM) with specified encoder parameters.
 
     Args:
         encoder_embed_dim (int | List[int]): Embedding dimension for the encoder.
         encoder_depth (int | List[int]): Depth of the encoder.
         encoder_num_heads (int | List[int]): Number of attention heads in the encoder.
         encoder_global_attn_indexes (List[int] | None): Indexes for global attention in the encoder.
-        checkpoint (str | None): Path to the model checkpoint file.
-        mobile_sam (bool): Whether to build a Mobile-SAM model.
+        checkpoint (str | None, optional): Path to the model checkpoint file.
+        mobile_sam (bool, optional): Whether to build a Mobile-SAM model.
 
     Returns:
         (SAMModel): A Segment Anything Model instance with the specified architecture.
@@ -224,17 +224,17 @@ def _build_sam2(
     checkpoint=None,
 ):
     """
-    Builds and returns a Segment Anything Model 2 (SAM2) with specified architecture parameters.
+    Build and return a Segment Anything Model 2 (SAM2) with specified architecture parameters.
 
     Args:
-        encoder_embed_dim (int): Embedding dimension for the encoder.
-        encoder_stages (List[int]): Number of blocks in each stage of the encoder.
-        encoder_num_heads (int): Number of attention heads in the encoder.
-        encoder_global_att_blocks (List[int]): Indices of global attention blocks in the encoder.
-        encoder_backbone_channel_list (List[int]): Channel dimensions for each level of the encoder backbone.
-        encoder_window_spatial_size (List[int]): Spatial size of the window for position embeddings.
-        encoder_window_spec (List[int]): Window specifications for each stage of the encoder.
-        checkpoint (str | None): Path to the checkpoint file for loading pre-trained weights.
+        encoder_embed_dim (int, optional): Embedding dimension for the encoder.
+        encoder_stages (List[int], optional): Number of blocks in each stage of the encoder.
+        encoder_num_heads (int, optional): Number of attention heads in the encoder.
+        encoder_global_att_blocks (List[int], optional): Indices of global attention blocks in the encoder.
+        encoder_backbone_channel_list (List[int], optional): Channel dimensions for each level of the encoder backbone.
+        encoder_window_spatial_size (List[int], optional): Spatial size of the window for position embeddings.
+        encoder_window_spec (List[int], optional): Window specifications for each stage of the encoder.
+        checkpoint (str | None, optional): Path to the checkpoint file for loading pre-trained weights.
 
     Returns:
         (SAM2Model): A configured and initialized SAM2 model.
@@ -326,10 +326,10 @@ sam_model_map = {
 
 def build_sam(ckpt="sam_b.pt"):
     """
-    Builds and returns a Segment Anything Model (SAM) based on the provided checkpoint.
+    Build and return a Segment Anything Model (SAM) based on the provided checkpoint.
 
     Args:
-        ckpt (str | Path): Path to the checkpoint file or name of a pre-defined SAM model.
+        ckpt (str | Path, optional): Path to the checkpoint file or name of a pre-defined SAM model.
 
     Returns:
         (SAMModel | SAM2Model): A configured and initialized SAM or SAM2 model instance.

ultralytics/models/sam/model.py

@@ -15,6 +15,7 @@ Key Features:
 """
 
 from pathlib import Path
+from typing import Dict, Type
 
 from ultralytics.engine.model import Model
 from ultralytics.utils.torch_utils import model_info
@@ -36,8 +37,8 @@ class SAM(Model):
         task (str): The task type, set to "segment" for SAM models.
 
     Methods:
-        predict: Performs segmentation prediction on the given image or video source.
-        info: Logs information about the SAM model.
+        predict: Perform segmentation prediction on the given image or video source.
+        info: Log information about the SAM model.
 
     Examples:
         >>> sam = SAM("sam_b.pt")
@@ -46,7 +47,7 @@ class SAM(Model):
         >>>     print(f"Detected {len(r.masks)} masks")
     """
 
-    def __init__(self, model="sam_b.pt") -> None:
+    def __init__(self, model: str = "sam_b.pt") -> None:
         """
         Initialize the SAM (Segment Anything Model) instance.
 
@@ -81,7 +82,7 @@ class SAM(Model):
 
         self.model = build_sam(weights)
 
-    def predict(self, source, stream=False, bboxes=None, points=None, labels=None, **kwargs):
+    def predict(self, source, stream: bool = False, bboxes=None, points=None, labels=None, **kwargs):
         """
         Perform segmentation prediction on the given image or video source.
 
@@ -108,7 +109,7 @@ class SAM(Model):
         prompts = dict(bboxes=bboxes, points=points, labels=labels)
         return super().predict(source, stream, prompts=prompts, **kwargs)
 
-    def __call__(self, source=None, stream=False, bboxes=None, points=None, labels=None, **kwargs):
+    def __call__(self, source=None, stream: bool = False, bboxes=None, points=None, labels=None, **kwargs):
         """
         Perform segmentation prediction on the given image or video source.
 
@@ -134,7 +135,7 @@ class SAM(Model):
         """
         return self.predict(source, stream, bboxes, points, labels, **kwargs)
 
-    def info(self, detailed=False, verbose=True):
+    def info(self, detailed: bool = False, verbose: bool = True):
         """
         Log information about the SAM model.
 
@@ -153,13 +154,13 @@ class SAM(Model):
         return model_info(self.model, detailed=detailed, verbose=verbose)
 
     @property
-    def task_map(self):
+    def task_map(self) -> Dict[str, Dict[str, Type[Predictor]]]:
         """
         Provide a mapping from the 'segment' task to its corresponding 'Predictor'.
 
         Returns:
-            (Dict[str, Dict[str, Type[Predictor]]]): A dictionary mapping the 'segment' task to its corresponding Predictor
-                class. For SAM2 models, it maps to SAM2Predictor, otherwise to the standard Predictor.
+            (Dict[str, Dict[str, Type[Predictor]]]): A dictionary mapping the 'segment' task to its corresponding
+                Predictor class. For SAM2 models, it maps to SAM2Predictor, otherwise to the standard Predictor.
 
         Examples:
             >>> sam = SAM("sam_b.pt")