dgenerate-ultralytics-headless 8.3.214__py3-none-any.whl → 8.4.7__py3-none-any.whl

ultralytics/models/rtdetr/val.py

@@ -16,8 +16,7 @@ __all__ = ("RTDETRValidator",)  # tuple or list
 
 
 class RTDETRDataset(YOLODataset):
-    """
-    Real-Time DEtection and TRacking (RT-DETR) dataset class extending the base YOLODataset class.
+    """Real-Time DEtection and TRacking (RT-DETR) dataset class extending the base YOLODataset class.
 
     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.
@@ -36,12 +35,11 @@ class RTDETRDataset(YOLODataset):
     Examples:
         Initialize an RT-DETR dataset
         >>> dataset = RTDETRDataset(img_path="path/to/images", imgsz=640)
-        >>> image, hw = dataset.load_image(0)
+        >>> image, hw0, hw = dataset.load_image(0)
     """
 
     def __init__(self, *args, data=None, **kwargs):
-        """
-        Initialize the RTDETRDataset class by inheriting from the YOLODataset class.
+        """Initialize the RTDETRDataset class by inheriting from the YOLODataset class.
 
         This constructor sets up a dataset specifically optimized for the RT-DETR (Real-Time DEtection and TRacking)
         model, building upon the base YOLODataset functionality.
@@ -54,27 +52,26 @@ class RTDETRDataset(YOLODataset):
         super().__init__(*args, data=data, **kwargs)
 
     def load_image(self, i, rect_mode=False):
-        """
-        Load one image from dataset index 'i'.
+        """Load one image from dataset index 'i'.
 
         Args:
             i (int): Index of the image to load.
             rect_mode (bool, optional): Whether to use rectangular mode for batch inference.
 
         Returns:
-            im (torch.Tensor): The loaded image.
-            resized_hw (tuple): Height and width of the resized image with shape (2,).
+            im (np.ndarray): Loaded image as a NumPy array.
+            hw_original (tuple[int, int]): Original image dimensions in (height, width) format.
+            hw_resized (tuple[int, int]): Resized image dimensions in (height, width) format.
 
         Examples:
             Load an image from the dataset
             >>> dataset = RTDETRDataset(img_path="path/to/images")
-            >>> image, hw = dataset.load_image(0)
+            >>> image, hw0, hw = dataset.load_image(0)
         """
         return super().load_image(i=i, rect_mode=rect_mode)
 
     def build_transforms(self, hyp=None):
-        """
-        Build transformation pipeline for the dataset.
+        """Build transformation pipeline for the dataset.
 
         Args:
             hyp (dict, optional): Hyperparameters for transformations.
@@ -105,8 +102,7 @@ class RTDETRDataset(YOLODataset):
 
 
 class RTDETRValidator(DetectionValidator):
-    """
-    RTDETRValidator extends the DetectionValidator class to provide validation capabilities specifically tailored for
+    """RTDETRValidator extends the DetectionValidator class to provide validation capabilities specifically tailored for
     the RT-DETR (Real-Time DETR) object detection model.
 
     The class allows building of an RTDETR-specific dataset for validation, applies Non-maximum suppression for
@@ -132,8 +128,7 @@ class RTDETRValidator(DetectionValidator):
     """
 
     def build_dataset(self, img_path, mode="val", batch=None):
-        """
-        Build an RTDETR Dataset.
+        """Build an RTDETR Dataset.
 
         Args:
             img_path (str): Path to the folder containing images.
@@ -156,15 +151,19 @@ class RTDETRValidator(DetectionValidator):
             data=self.data,
         )
 
+    def scale_preds(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> dict[str, torch.Tensor]:
+        """Scales predictions to the original image size."""
+        return predn
+
     def postprocess(
         self, preds: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor]
     ) -> list[dict[str, torch.Tensor]]:
-        """
-        Apply Non-maximum suppression to prediction outputs.
+        """Apply Non-maximum suppression to prediction outputs.
 
         Args:
             preds (torch.Tensor | list | tuple): Raw predictions from the model. If tensor, should have shape
-                (batch_size, num_predictions, num_classes + 4) where last dimension contains bbox coords and class scores.
+                (batch_size, num_predictions, num_classes + 4) where last dimension contains bbox coords and
+                class scores.
 
         Returns:
             (list[dict[str, torch.Tensor]]): List of dictionaries for each image, each containing:
@@ -190,12 +189,11 @@ class RTDETRValidator(DetectionValidator):
         return [{"bboxes": x[:, :4], "conf": x[:, 4], "cls": x[:, 5]} for x in outputs]
 
     def pred_to_json(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> None:
-        """
-        Serialize YOLO predictions to COCO json format.
+        """Serialize YOLO predictions to COCO json format.
 
         Args:
-            predn (dict[str, torch.Tensor]): Predictions dictionary containing 'bboxes', 'conf', and 'cls' keys
-                with bounding box coordinates, confidence scores, and class predictions.
+            predn (dict[str, torch.Tensor]): Predictions dictionary containing 'bboxes', 'conf', and 'cls' keys with
+                bounding box coordinates, confidence scores, and class predictions.
             pbatch (dict[str, Any]): Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'.
         """
         path = Path(pbatch["im_file"])

ultralytics/models/sam/__init__.py

@@ -1,12 +1,25 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from .model import SAM
-from .predict import Predictor, SAM2DynamicInteractivePredictor, SAM2Predictor, SAM2VideoPredictor
+from .predict import (
+    Predictor,
+    SAM2DynamicInteractivePredictor,
+    SAM2Predictor,
+    SAM2VideoPredictor,
+    SAM3Predictor,
+    SAM3SemanticPredictor,
+    SAM3VideoPredictor,
+    SAM3VideoSemanticPredictor,
+)
 
 __all__ = (
     "SAM",
     "Predictor",
+    "SAM2DynamicInteractivePredictor",
     "SAM2Predictor",
     "SAM2VideoPredictor",
-    "SAM2DynamicInteractivePredictor",
+    "SAM3Predictor",
+    "SAM3SemanticPredictor",
+    "SAM3VideoPredictor",
+    "SAM3VideoSemanticPredictor",
 )  # tuple or list of exportable items

ultralytics/models/sam/amg.py

@@ -14,8 +14,7 @@ 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:
-    """
-    Determine 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.
@@ -42,8 +41,7 @@ def is_box_near_crop_edge(
 
 
 def batch_iterator(batch_size: int, *args) -> Generator[list[Any]]:
-    """
-    Yield batches of data from input arguments with specified batch size for efficient processing.
+    """Yield batches of data from input arguments with specified batch size for efficient processing.
 
     This function takes a batch size and any number of iterables, then yields batches of elements from those
     iterables. All input iterables must have the same length.
@@ -71,11 +69,10 @@ def batch_iterator(batch_size: int, *args) -> Generator[list[Any]]:
 
 
 def calculate_stability_score(masks: torch.Tensor, mask_threshold: float, threshold_offset: float) -> torch.Tensor:
-    """
-    Compute 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.
+    The stability score is the IoU between binary masks obtained by thresholding the predicted mask logits at high and
+    low values.
 
     Args:
         masks (torch.Tensor): Batch of predicted mask logits.
@@ -85,15 +82,15 @@ def calculate_stability_score(masks: torch.Tensor, mask_threshold: float, thresh
     Returns:
         (torch.Tensor): Stability scores for each mask in the batch.
 
-    Notes:
-        - One mask is always contained inside the other.
-        - Memory is saved by preventing unnecessary cast to torch.int64.
-
     Examples:
         >>> masks = torch.rand(10, 256, 256)  # Batch of 10 masks
         >>> mask_threshold = 0.5
         >>> threshold_offset = 0.1
         >>> stability_scores = calculate_stability_score(masks, mask_threshold, threshold_offset)
+
+    Notes:
+        - One mask is always contained inside the other.
+        - Memory is saved by preventing unnecessary cast to torch.int64.
     """
     intersections = (masks > (mask_threshold + threshold_offset)).sum(-1, dtype=torch.int16).sum(-1, dtype=torch.int32)
     unions = (masks > (mask_threshold - threshold_offset)).sum(-1, dtype=torch.int16).sum(-1, dtype=torch.int32)
@@ -117,8 +114,7 @@ def build_all_layer_point_grids(n_per_side: int, n_layers: int, scale_per_layer:
 def generate_crop_boxes(
     im_size: tuple[int, ...], n_layers: int, overlap_ratio: float
 ) -> tuple[list[list[int]], list[int]]:
-    """
-    Generate 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.
@@ -145,7 +141,7 @@ def generate_crop_boxes(
 
     def crop_len(orig_len, n_crops, 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))
+        return math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops)
 
     for i_layer in range(n_layers):
         n_crops_per_side = 2 ** (i_layer + 1)
@@ -198,8 +194,7 @@ 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]:
-    """
-    Remove 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.
@@ -227,7 +222,7 @@ def remove_small_regions(mask: np.ndarray, area_thresh: float, mode: str) -> tup
     small_regions = [i + 1 for i, s in enumerate(sizes) if s < area_thresh]
     if not small_regions:
         return mask, False
-    fill_labels = [0] + small_regions
+    fill_labels = [0, *small_regions]
     if not correct_holes:
         # If every region is below threshold, keep largest
         fill_labels = [i for i in range(n_labels) if i not in fill_labels] or [int(np.argmax(sizes)) + 1]
@@ -236,8 +231,7 @@ def remove_small_regions(mask: np.ndarray, area_thresh: float, mode: str) -> tup
 
 
 def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor:
-    """
-    Calculate 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

@@ -11,6 +11,7 @@ from functools import partial
 import torch
 
 from ultralytics.utils.downloads import attempt_download_asset
+from ultralytics.utils.patches import torch_load
 
 from .modules.decoders import MaskDecoder
 from .modules.encoders import FpnNeck, Hiera, ImageEncoder, ImageEncoderViT, MemoryEncoder, PromptEncoder
@@ -20,6 +21,21 @@ from .modules.tiny_encoder import TinyViT
 from .modules.transformer import TwoWayTransformer
 
 
+def _load_checkpoint(model, checkpoint):
+    """Load checkpoint into model from file path."""
+    if checkpoint is None:
+        return model
+
+    checkpoint = attempt_download_asset(checkpoint)
+    with open(checkpoint, "rb") as f:
+        state_dict = torch_load(f)
+    # Handle nested "model" key
+    if "model" in state_dict and isinstance(state_dict["model"], dict):
+        state_dict = state_dict["model"]
+    model.load_state_dict(state_dict)
+    return model
+
+
 def build_sam_vit_h(checkpoint=None):
     """Build and return a Segment Anything Model (SAM) h-size model with specified encoder parameters."""
     return _build_sam(
@@ -126,8 +142,7 @@ def _build_sam(
     checkpoint=None,
     mobile_sam=False,
 ):
-    """
-    Build 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.
@@ -205,26 +220,22 @@ def _build_sam(
         pixel_std=[58.395, 57.12, 57.375],
     )
     if checkpoint is not None:
-        checkpoint = attempt_download_asset(checkpoint)
-        with open(checkpoint, "rb") as f:
-            state_dict = torch.load(f)
-        sam.load_state_dict(state_dict)
+        sam = _load_checkpoint(sam, checkpoint)
     sam.eval()
     return sam
 
 
 def _build_sam2(
     encoder_embed_dim=1280,
-    encoder_stages=[2, 6, 36, 4],
+    encoder_stages=(2, 6, 36, 4),
     encoder_num_heads=2,
-    encoder_global_att_blocks=[7, 15, 23, 31],
-    encoder_backbone_channel_list=[1152, 576, 288, 144],
-    encoder_window_spatial_size=[7, 7],
-    encoder_window_spec=[8, 4, 16, 8],
+    encoder_global_att_blocks=(7, 15, 23, 31),
+    encoder_backbone_channel_list=(1152, 576, 288, 144),
+    encoder_window_spatial_size=(7, 7),
+    encoder_window_spec=(8, 4, 16, 8),
     checkpoint=None,
 ):
-    """
-    Build and return 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, optional): Embedding dimension for the encoder.
@@ -300,10 +311,7 @@ def _build_sam2(
     )
 
     if checkpoint is not None:
-        checkpoint = attempt_download_asset(checkpoint)
-        with open(checkpoint, "rb") as f:
-            state_dict = torch.load(f)["model"]
-        sam2.load_state_dict(state_dict)
+        sam2 = _load_checkpoint(sam2, checkpoint)
     sam2.eval()
     return sam2
 
@@ -325,8 +333,7 @@ sam_model_map = {
 
 
 def build_sam(ckpt="sam_b.pt"):
-    """
-    Build and return 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, optional): Path to the checkpoint file or name of a pre-defined SAM model.