ultralytics 8.3.101__tar.gz → 8.3.102__tar.gz

{ultralytics-8.3.101/ultralytics.egg-info → ultralytics-8.3.102}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ultralytics
-Version: 8.3.101
+Version: 8.3.102
 Summary: Ultralytics YOLO 🚀 for SOTA object detection, multi-object tracking, instance segmentation, pose estimation and image classification.
 Author-email: Glenn Jocher <glenn.jocher@ultralytics.com>, Jing Qiu <jing.qiu@ultralytics.com>
 Maintainer-email: Ultralytics <hello@ultralytics.com>
@@ -56,7 +56,6 @@ Requires-Dist: coverage[toml]; extra == "dev"
 Requires-Dist: mkdocs>=1.6.0; extra == "dev"
 Requires-Dist: mkdocs-material>=9.5.9; extra == "dev"
 Requires-Dist: mkdocstrings[python]; extra == "dev"
-Requires-Dist: mkdocs-redirects; extra == "dev"
 Requires-Dist: mkdocs-ultralytics-plugin>=0.1.17; extra == "dev"
 Requires-Dist: mkdocs-macros-plugin>=1.0.5; extra == "dev"
 Provides-Extra: export

{ultralytics-8.3.101 → ultralytics-8.3.102}/pyproject.toml

@@ -89,7 +89,6 @@ dev = [
     "mkdocs>=1.6.0",
     "mkdocs-material>=9.5.9",
     "mkdocstrings[python]",
-    "mkdocs-redirects", # 301 redirects
     "mkdocs-ultralytics-plugin>=0.1.17", # for meta descriptions and images, dates and authors
     "mkdocs-macros-plugin>=1.0.5"  # duplicating content (i.e. export tables) in multiple places
 ]

ultralytics-8.3.102/tests/test_solutions.py

@@ -0,0 +1,167 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+# Tests Ultralytics Solutions: https://docs.ultralytics.com/solutions/,
+# including every solution excluding DistanceCalculation and Security Alarm System.
+
+import cv2
+import pytest
+
+from tests import TMP
+from ultralytics import solutions
+from ultralytics.utils import ASSETS_URL
+from ultralytics.utils.downloads import safe_download
+
+# Pre-defined arguments values
+SHOW = False
+DEMO_VIDEO = "solutions_ci_demo.mp4"  # for all the solutions, except workout, object cropping and parking management
+CROP_VIDEO = "decelera_landscape_min.mov"  # for object cropping solution
+POSE_VIDEO = "solution_ci_pose_demo.mp4"  # only for workouts monitoring solution
+PARKING_VIDEO = "solution_ci_parking_demo.mp4"  # only for parking management solution
+PARKING_AREAS_JSON = "solution_ci_parking_areas.json"  # only for parking management solution
+PARKING_MODEL = "solutions_ci_parking_model.pt"  # only for parking management solution
+MODEL_FILE = "yolo11n.pt"  # model file used for solutions, except parking management and instance segmentation
+REGION = [(20, 400), (1080, 400), (1080, 360), (20, 360)]  # for object counting, speed estimation and queue management
+
+# Test configs for each solution : (name, class, needs_frame_count, video, kwargs)
+SOLUTIONS = [
+    (
+        "ObjectCounter",
+        solutions.ObjectCounter,
+        False,
+        DEMO_VIDEO,
+        {"region": REGION, "model": MODEL_FILE, "show": SHOW},
+    ),
+    (
+        "Heatmap",
+        solutions.Heatmap,
+        False,
+        DEMO_VIDEO,
+        {"colormap": cv2.COLORMAP_PARULA, "model": MODEL_FILE, "show": SHOW, "region": None},
+    ),
+    (
+        "HeatmapWithRegion",
+        solutions.Heatmap,
+        False,
+        DEMO_VIDEO,
+        {"colormap": cv2.COLORMAP_PARULA, "region": REGION, "model": MODEL_FILE, "show": SHOW},
+    ),
+    (
+        "SpeedEstimator",
+        solutions.SpeedEstimator,
+        False,
+        DEMO_VIDEO,
+        {"region": REGION, "model": MODEL_FILE, "show": SHOW},
+    ),
+    (
+        "QueueManager",
+        solutions.QueueManager,
+        False,
+        DEMO_VIDEO,
+        {"region": REGION, "model": MODEL_FILE, "show": SHOW},
+    ),
+    (
+        "LineAnalytics",
+        solutions.Analytics,
+        True,
+        DEMO_VIDEO,
+        {"analytics_type": "line", "model": MODEL_FILE, "show": SHOW},
+    ),
+    (
+        "PieAnalytics",
+        solutions.Analytics,
+        True,
+        DEMO_VIDEO,
+        {"analytics_type": "pie", "model": MODEL_FILE, "show": SHOW},
+    ),
+    (
+        "BarAnalytics",
+        solutions.Analytics,
+        True,
+        DEMO_VIDEO,
+        {"analytics_type": "bar", "model": MODEL_FILE, "show": SHOW},
+    ),
+    (
+        "AreaAnalytics",
+        solutions.Analytics,
+        True,
+        DEMO_VIDEO,
+        {"analytics_type": "area", "model": MODEL_FILE, "show": SHOW},
+    ),
+    ("TrackZone", solutions.TrackZone, False, DEMO_VIDEO, {"region": REGION, "model": MODEL_FILE, "show": SHOW}),
+    (
+        "ObjectCropper",
+        solutions.ObjectCropper,
+        False,
+        CROP_VIDEO,
+        {"crop_dir": str(TMP / "cropped-detections"), "model": MODEL_FILE, "show": SHOW},
+    ),
+    (
+        "ObjectBlurrer",
+        solutions.ObjectBlurrer,
+        False,
+        DEMO_VIDEO,
+        {"blur_ratio": 0.5, "model": MODEL_FILE, "show": SHOW},
+    ),
+    (
+        "InstanceSegmentation",
+        solutions.InstanceSegmentation,
+        False,
+        DEMO_VIDEO,
+        {"model": "yolo11n-seg.pt", "show": SHOW},
+    ),
+    ("VisionEye", solutions.VisionEye, False, DEMO_VIDEO, {"model": MODEL_FILE, "show": SHOW}),
+    (
+        "RegionCounter",
+        solutions.RegionCounter,
+        False,
+        DEMO_VIDEO,
+        {"region": REGION, "model": MODEL_FILE, "show": SHOW},
+    ),
+    ("AIGym", solutions.AIGym, False, POSE_VIDEO, {"kpts": [6, 8, 10], "show": SHOW}),
+    (
+        "ParkingManager",
+        solutions.ParkingManagement,
+        False,
+        PARKING_VIDEO,
+        {"model": str(TMP / PARKING_MODEL), "show": SHOW, "json_file": str(TMP / PARKING_AREAS_JSON)},
+    ),
+    (
+        "StreamlitInference",
+        solutions.Inference,
+        False,
+        None,  # streamlit application don't require video file
+        {},  # streamlit application don't accept arguments
+    ),
+]
+
+
+def process_video(solution, video_path, needs_frame_count=False):
+    """Process video with solution, feeding frames and optional frame count."""
+    cap = cv2.VideoCapture(video_path)
+    assert cap.isOpened(), f"Error reading video file {video_path}"
+
+    frame_count = 0
+    while cap.isOpened():
+        success, im0 = cap.read()
+        if not success:
+            break
+        frame_count += 1
+        im_copy = im0.copy()
+        args = [im_copy, frame_count] if needs_frame_count else [im_copy]
+        _ = solution(*args)
+
+    cap.release()
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("name, solution_class, needs_frame_count, video, kwargs", SOLUTIONS)
+def test_solution(name, solution_class, needs_frame_count, video, kwargs):
+    """Test individual Ultralytics solution."""
+    safe_download(url=f"{ASSETS_URL}/{video}", dir=TMP)
+    if name == "ParkingManager":
+        safe_download(url=f"{ASSETS_URL}/{PARKING_AREAS_JSON}", dir=TMP)
+        safe_download(url=f"{ASSETS_URL}/{PARKING_MODEL}", dir=TMP)
+    solution = solution_class(**kwargs)
+    solution.inference() if name == "StreamlitInference" else process_video(
+        solution, str(TMP / video), needs_frame_count
+    )

{ultralytics-8.3.101 → ultralytics-8.3.102}/ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-__version__ = "8.3.101"
+__version__ = "8.3.102"
 
 import os
 

{ultralytics-8.3.101 → ultralytics-8.3.102}/ultralytics/engine/exporter.py

@@ -115,18 +115,32 @@ def export_formats():
         ["PyTorch", "-", ".pt", True, True, []],
         ["TorchScript", "torchscript", ".torchscript", True, True, ["batch", "optimize", "nms"]],
         ["ONNX", "onnx", ".onnx", True, True, ["batch", "dynamic", "half", "opset", "simplify", "nms"]],
-        ["OpenVINO", "openvino", "_openvino_model", True, False, ["batch", "dynamic", "half", "int8", "nms"]],
-        ["TensorRT", "engine", ".engine", False, True, ["batch", "dynamic", "half", "int8", "simplify", "nms"]],
+        [
+            "OpenVINO",
+            "openvino",
+            "_openvino_model",
+            True,
+            False,
+            ["batch", "dynamic", "half", "int8", "nms", "fraction"],
+        ],
+        [
+            "TensorRT",
+            "engine",
+            ".engine",
+            False,
+            True,
+            ["batch", "dynamic", "half", "int8", "simplify", "nms", "fraction"],
+        ],
         ["CoreML", "coreml", ".mlpackage", True, False, ["batch", "half", "int8", "nms"]],
         ["TensorFlow SavedModel", "saved_model", "_saved_model", True, True, ["batch", "int8", "keras", "nms"]],
         ["TensorFlow GraphDef", "pb", ".pb", True, True, ["batch"]],
-        ["TensorFlow Lite", "tflite", ".tflite", True, False, ["batch", "half", "int8", "nms"]],
+        ["TensorFlow Lite", "tflite", ".tflite", True, False, ["batch", "half", "int8", "nms", "fraction"]],
         ["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", True, False, []],
         ["TensorFlow.js", "tfjs", "_web_model", True, False, ["batch", "half", "int8", "nms"]],
         ["PaddlePaddle", "paddle", "_paddle_model", True, True, ["batch"]],
         ["MNN", "mnn", ".mnn", True, True, ["batch", "half", "int8"]],
         ["NCNN", "ncnn", "_ncnn_model", True, True, ["batch", "half"]],
-        ["IMX", "imx", "_imx_model", True, True, ["int8"]],
+        ["IMX", "imx", "_imx_model", True, True, ["int8", "fraction"]],
         ["RKNN", "rknn", "_rknn_model", False, False, ["batch", "name"]],
     ]
     return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU", "Arguments"], zip(*x)))
@@ -144,7 +158,7 @@ def validate_args(format, passed_args, valid_args):
     Raises:
         AssertionError: If an unsupported argument is used, or if the format lacks supported argument listings.
     """
-    export_args = ["half", "int8", "dynamic", "keras", "nms", "batch"]
+    export_args = ["half", "int8", "dynamic", "keras", "nms", "batch", "fraction"]
 
     assert valid_args is not None, f"ERROR ❌️ valid arguments for '{format}' not listed."
     custom = {"batch": 1, "data": None, "device": None}  # exporter defaults
@@ -493,6 +507,7 @@ class Exporter:
         dataset = YOLODataset(
             data[self.args.split or "val"],
             data=data,
+            fraction=self.args.fraction,
             task=self.model.task,
             imgsz=self.imgsz[0],
             augment=False,

{ultralytics-8.3.101 → ultralytics-8.3.102}/ultralytics/hub/__init__.py

@@ -96,7 +96,21 @@ def export_fmts_hub():
 
 
 def export_model(model_id: str = "", format: str = "torchscript"):
-    """Export a model to the specified format."""
+    """
+    Export a model to a specified format for deployment via the Ultralytics HUB API.
+
+    Args:
+        model_id (str): The ID of the model to export. An empty string will use the default model.
+        format (str): The format to export the model to. Must be one of the supported formats returned by
+            export_fmts_hub().
+
+    Raises:
+        AssertionError: If the specified format is not supported or if the export request fails.
+
+    Examples:
+        >>> from ultralytics import hub
+        >>> hub.export_model(model_id="your_model_id", format="torchscript")
+    """
     assert format in export_fmts_hub(), f"Unsupported export format '{format}', valid formats are {export_fmts_hub()}"
     r = requests.post(
         f"{HUB_API_ROOT}/v1/models/{model_id}/export", json={"format": format}, headers={"x-api-key": Auth().api_key}
@@ -106,7 +120,20 @@ def export_model(model_id: str = "", format: str = "torchscript"):
 
 
 def get_export(model_id: str = "", format: str = "torchscript"):
-    """Get an exported model dictionary with download URL."""
+    """
+    Retrieve an exported model in the specified format from Ultralytics HUB using the model ID.
+
+    Args:
+        model_id (str): The ID of the model to retrieve from Ultralytics HUB.
+        format (str): The export format to retrieve. Must be one of the supported formats returned by export_fmts_hub().
+
+    Raises:
+        AssertionError: If the specified format is not supported or if the API request fails.
+
+    Examples:
+        >>> from ultralytics import hub
+        >>> hub.get_export(model_id="your_model_id", format="torchscript")
+    """
     assert format in export_fmts_hub(), f"Unsupported export format '{format}', valid formats are {export_fmts_hub()}"
     r = requests.post(
         f"{HUB_API_ROOT}/get-export",

{ultralytics-8.3.101 → ultralytics-8.3.102}/ultralytics/hub/google/__init__.py

@@ -83,7 +83,24 @@ class GCPRegions:
 
     @staticmethod
     def _ping_region(region: str, attempts: int = 1) -> Tuple[str, float, float, float, float]:
-        """Pings a specified GCP region and returns latency statistics: mean, min, max, and standard deviation."""
+        """
+        Ping a specified GCP region and measure network latency statistics.
+
+        Args:
+                region (str): The GCP region identifier to ping (e.g., 'us-central1').
+                attempts (int): Number of ping attempts to make for calculating statistics.
+
+        Returns:
+                region (str): The GCP region identifier that was pinged.
+                mean_latency (float): Mean latency in milliseconds, or infinity if all pings failed.
+                min_latency (float): Minimum latency in milliseconds, or infinity if all pings failed.
+                max_latency (float): Maximum latency in milliseconds, or infinity if all pings failed.
+                std_dev (float): Standard deviation of latencies in milliseconds, or infinity if all pings failed.
+
+        Examples:
+                >>> region, mean, min_lat, max_lat, std = GCPRegions._ping_region("us-central1", attempts=3)
+                >>> print(f"Region {region} has mean latency: {mean:.2f}ms")
+        """
         url = f"https://{region}-docker.pkg.dev"
         latencies = []
         for _ in range(attempts):

{ultralytics-8.3.101 → ultralytics-8.3.102}/ultralytics/models/fastsam/predict.py

@@ -33,7 +33,18 @@ class FastSAMPredictor(SegmentationPredictor):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """Initialize the FastSAMPredictor with configuration and callbacks."""
+        """
+        Initialize the FastSAMPredictor with configuration and callbacks.
+
+        This initializes a predictor specialized for Fast SAM (Segment Anything Model) segmentation tasks. The predictor
+        extends SegmentationPredictor with custom post-processing for mask prediction and non-maximum suppression
+        optimized for single-class segmentation.
+
+        Args:
+            cfg (dict): Configuration for the predictor. Defaults to Ultralytics DEFAULT_CFG.
+            overrides (dict, optional): Configuration overrides.
+            _callbacks (list, optional): List of callback functions.
+        """
         super().__init__(cfg, overrides, _callbacks)
         self.prompts = {}
 

{ultralytics-8.3.101 → ultralytics-8.3.102}/ultralytics/models/nas/predict.py

@@ -33,8 +33,26 @@ class NASPredictor(DetectionPredictor):
     """
 
     def postprocess(self, preds_in, img, orig_imgs):
-        """Postprocess predictions and returns a list of Results objects."""
-        # Convert boxes from xyxy to xywh format and concatenate with class scores
+        """
+        Postprocess NAS model predictions to generate final detection results.
+
+        This method takes raw predictions from a YOLO NAS model, converts bounding box formats, and applies
+        post-processing operations to generate the final detection results compatible with Ultralytics
+        result visualization and analysis tools.
+
+        Args:
+            preds_in (list): Raw predictions from the NAS model, typically containing bounding boxes and class scores.
+            img (torch.Tensor): Input image tensor that was fed to the model, with shape (B, C, H, W).
+            orig_imgs (list | torch.Tensor | np.ndarray): Original images before preprocessing, used for scaling
+                coordinates back to original dimensions.
+
+        Returns:
+            (list): List of Results objects containing the processed predictions for each image in the batch.
+
+        Examples:
+            >>> predictor = NAS("yolo_nas_s").predictor
+            >>> results = predictor.postprocess(raw_preds, img, orig_imgs)
+        """
         boxes = ops.xyxy2xywh(preds_in[0][0])
-        preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)
+        preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)  # concatenate with class scores
         return super().postprocess(preds, img, orig_imgs)

{ultralytics-8.3.101 → ultralytics-8.3.102}/ultralytics/models/rtdetr/val.py

@@ -19,11 +19,35 @@ class RTDETRDataset(YOLODataset):
     """
 
     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.
+
+        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.
+            **kwargs (Any): Additional keyword arguments passed to the parent YOLODataset class.
+        """
         super().__init__(*args, data=data, **kwargs)
 
     def load_image(self, i, rect_mode=False):
-        """Loads 1 image from dataset index 'i', returns (im, resized hw)."""
+        """
+        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 (numpy.ndarray): The loaded image.
+            resized_hw (tuple): Height and width of the resized image with shape (2,).
+
+        Examples:
+            >>> dataset = RTDETRDataset(...)
+            >>> image, hw = dataset.load_image(0)
+        """
         return super().load_image(i=i, rect_mode=rect_mode)
 
     def build_transforms(self, hyp=None):

{ultralytics-8.3.101 → ultralytics-8.3.102}/ultralytics/models/sam/amg.py

@@ -22,7 +22,28 @@ def is_box_near_crop_edge(
 
 
 def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]:
-    """Yields 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.
+
+    Args:
+        batch_size (int): Size of each batch to yield.
+        *args (Any): Variable length input iterables to batch. All iterables must have the same length.
+
+    Yields:
+        (List[Any]): A list of batched elements from each input iterable.
+
+    Examples:
+        >>> data = [1, 2, 3, 4, 5]
+        >>> labels = ["a", "b", "c", "d", "e"]
+        >>> for batch in batch_iterator(2, data, labels):
+        ...     print(batch)
+        [[1, 2], ['a', 'b']]
+        [[3, 4], ['c', 'd']]
+        [[5], ['e']]
+    """
     assert args and all(len(a) == len(args[0]) for a in args), "Batched iteration must have same-size inputs."
     n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0)
     for b in range(n_batches):

{ultralytics-8.3.101 → ultralytics-8.3.102}/ultralytics/models/sam/modules/encoders.py

@@ -386,7 +386,24 @@ class MemoryEncoder(nn.Module):
         out_dim,
         in_dim=256,  # in_dim of pix_feats
     ):
-        """Initialize the MemoryEncoder for encoding pixel features and masks into memory representations."""
+        """
+        Initialize the MemoryEncoder for encoding pixel features and masks into memory representations.
+
+        This encoder processes pixel-level features and masks, fusing them to generate encoded memory representations
+        suitable for downstream tasks in image segmentation models like SAM (Segment Anything Model).
+
+        Args:
+            out_dim (int): Output dimension of the encoded features.
+            in_dim (int): Input dimension of the pixel features. Default is 256.
+
+        Examples:
+            >>> encoder = MemoryEncoder(out_dim=256, in_dim=256)
+            >>> pix_feat = torch.randn(1, 256, 64, 64)
+            >>> masks = torch.randn(1, 1, 64, 64)
+            >>> encoded_feat, pos = encoder(pix_feat, masks)
+            >>> print(encoded_feat.shape, pos.shape)
+            torch.Size([1, 256, 64, 64]) torch.Size([1, 128, 64, 64])
+        """
         super().__init__()
 
         self.mask_downsampler = MaskDownSampler(kernel_size=3, stride=2, padding=1)
@@ -453,7 +470,26 @@ class ImageEncoder(nn.Module):
         neck: nn.Module,
         scalp: int = 0,
     ):
-        """Initialize the ImageEncoder with trunk and neck networks for feature extraction and refinement."""
+        """
+        Initialize the ImageEncoder with trunk and neck networks for feature extraction and refinement.
+
+        This encoder combines a trunk network for feature extraction with a neck network for feature refinement
+        and positional encoding generation. It can optionally discard the lowest resolution features.
+
+        Args:
+            trunk (nn.Module): The trunk network for initial feature extraction.
+            neck (nn.Module): The neck network for feature refinement and positional encoding generation.
+            scalp (int): Number of lowest resolution feature levels to discard.
+
+        Examples:
+            >>> trunk = SomeTrunkNetwork()
+            >>> neck = SomeNeckNetwork()
+            >>> encoder = ImageEncoder(trunk, neck, scalp=1)
+            >>> image = torch.randn(1, 3, 224, 224)
+            >>> output = encoder(image)
+            >>> print(output.keys())
+            dict_keys(['vision_features', 'vision_pos_enc', 'backbone_fpn'])
+        """
         super().__init__()
         self.trunk = trunk
         self.neck = neck
@@ -681,7 +717,34 @@ class Hiera(nn.Module):
         ),
         return_interm_layers=True,  # return feats from every stage
     ):
-        """Initialize the Hiera model, configuring its hierarchical vision transformer architecture."""
+        """
+        Initialize a Hiera model, a hierarchical vision transformer for efficient multiscale feature extraction.
+
+        Hiera is a hierarchical vision transformer architecture designed for efficient multiscale feature extraction
+        in image processing tasks. It uses a series of transformer blocks organized into stages, with optional
+        pooling and global attention mechanisms.
+
+        Args:
+            embed_dim (int): Initial embedding dimension for the model.
+            num_heads (int): Initial number of attention heads.
+            drop_path_rate (float): Stochastic depth rate.
+            q_pool (int): Number of query pooling stages.
+            q_stride (Tuple[int, int]): Downsampling stride between stages.
+            stages (Tuple[int, ...]): Number of blocks per stage.
+            dim_mul (float): Dimension multiplier factor at stage transitions.
+            head_mul (float): Head multiplier factor at stage transitions.
+            window_pos_embed_bkg_spatial_size (Tuple[int, int]): Spatial size for window positional embedding background.
+            window_spec (Tuple[int, ...]): Window sizes for each stage when not using global attention.
+            global_att_blocks (Tuple[int, ...]): Indices of blocks that use global attention.
+            return_interm_layers (bool): Whether to return intermediate layer outputs.
+
+        Examples:
+            >>> model = Hiera(embed_dim=96, num_heads=1, stages=(2, 3, 16, 3))
+            >>> input_tensor = torch.randn(1, 3, 224, 224)
+            >>> output_features = model(input_tensor)
+            >>> for feat in output_features:
+            ...     print(feat.shape)
+        """
         super().__init__()
 
         assert len(stages) == len(window_spec)
@@ -756,7 +819,25 @@ class Hiera(nn.Module):
         return pos_embed
 
     def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
-        """Perform forward pass through Hiera model, extracting multiscale features from input images."""
+        """
+        Perform forward pass through Hiera model, extracting multiscale features from input images.
+
+        Args:
+            x (torch.Tensor): Input tensor with shape (B, C, H, W) representing a batch of images.
+
+        Returns:
+            (List[torch.Tensor]): List of feature maps at different scales, each with shape (B, C_i, H_i, W_i), where
+                C_i is the channel dimension and H_i, W_i are the spatial dimensions at scale i. The list is ordered
+                from highest resolution (fine features) to lowest resolution (coarse features) if return_interm_layers
+                is True, otherwise contains only the final output.
+
+        Examples:
+            >>> model = Hiera(embed_dim=96, num_heads=1, stages=(2, 3, 16, 3))
+            >>> input_tensor = torch.randn(1, 3, 224, 224)
+            >>> output_features = model(input_tensor)
+            >>> for feat in output_features:
+            ...     print(feat.shape)
+        """
         x = self.patch_embed(x)
         # x: (B, H, W, C)
 

{ultralytics-8.3.101 → ultralytics-8.3.102}/ultralytics/models/sam/modules/memory_attention.py

@@ -60,7 +60,17 @@ class MemoryAttentionLayer(nn.Module):
         pos_enc_at_cross_attn_keys: bool = True,
         pos_enc_at_cross_attn_queries: bool = False,
     ):
-        """Initialize a memory attention layer with self-attention, cross-attention, and feedforward components."""
+        """
+        Initialize a memory attention layer with self-attention, cross-attention, and feedforward components.
+
+        Args:
+            d_model (int): Dimensionality of the model.
+            dim_feedforward (int): Dimensionality of the feedforward network.
+            dropout (float): Dropout rate for regularization.
+            pos_enc_at_attn (bool): Whether to add positional encoding at attention.
+            pos_enc_at_cross_attn_keys (bool): Whether to add positional encoding to cross-attention keys.
+            pos_enc_at_cross_attn_queries (bool): Whether to add positional encoding to cross-attention queries.
+        """
         super().__init__()
         self.d_model = d_model
         self.dim_feedforward = dim_feedforward
@@ -183,7 +193,31 @@ class MemoryAttention(nn.Module):
         num_layers: int,
         batch_first: bool = True,  # Do layers expect batch first input?
     ):
-        """Initialize MemoryAttention with specified layers and normalization for sequential data processing."""
+        """
+        Initialize MemoryAttention with specified layers and normalization for sequential data processing.
+
+        This class implements a multi-layer attention mechanism that combines self-attention and cross-attention
+        for processing sequential data, particularly useful in transformer-like architectures.
+
+        Args:
+            d_model (int): The dimension of the model's hidden state.
+            pos_enc_at_input (bool): Whether to apply positional encoding at the input.
+            layer (nn.Module): The attention layer to be used in the module.
+            num_layers (int): The number of attention layers.
+            batch_first (bool): Whether the input tensors are in batch-first format.
+
+        Examples:
+            >>> d_model = 256
+            >>> layer = MemoryAttentionLayer(d_model)
+            >>> attention = MemoryAttention(d_model, pos_enc_at_input=True, layer=layer, num_layers=3)
+            >>> curr = torch.randn(10, 32, d_model)  # (seq_len, batch_size, d_model)
+            >>> memory = torch.randn(20, 32, d_model)  # (mem_len, batch_size, d_model)
+            >>> curr_pos = torch.randn(10, 32, d_model)
+            >>> memory_pos = torch.randn(20, 32, d_model)
+            >>> output = attention(curr, memory, curr_pos, memory_pos)
+            >>> print(output.shape)
+            torch.Size([10, 32, 256])
+        """
         super().__init__()
         self.d_model = d_model
         self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(num_layers)])
@@ -200,7 +234,31 @@ class MemoryAttention(nn.Module):
         memory_pos: Optional[Tensor] = None,  # pos_enc for cross-attention inputs
         num_obj_ptr_tokens: int = 0,  # number of object pointer *tokens*
     ) -> torch.Tensor:
-        """Process inputs through attention layers, applying self and cross-attention with positional encoding."""
+        """
+        Process inputs through attention layers, applying self and cross-attention with positional encoding.
+
+        Args:
+            curr (torch.Tensor): Self-attention input tensor, representing the current state.
+            memory (torch.Tensor): Cross-attention input tensor, representing memory information.
+            curr_pos (Optional[Tensor]): Positional encoding for self-attention inputs.
+            memory_pos (Optional[Tensor]): Positional encoding for cross-attention inputs.
+            num_obj_ptr_tokens (int): Number of object pointer tokens to exclude from rotary position embedding.
+
+        Returns:
+            (torch.Tensor): Processed output tensor after applying attention layers and normalization.
+
+        Examples:
+            >>> d_model = 256
+            >>> layer = MemoryAttentionLayer(d_model)
+            >>> attention = MemoryAttention(d_model, pos_enc_at_input=True, layer=layer, num_layers=3)
+            >>> curr = torch.randn(10, 32, d_model)  # (seq_len, batch_size, d_model)
+            >>> memory = torch.randn(20, 32, d_model)  # (mem_len, batch_size, d_model)
+            >>> curr_pos = torch.randn(10, 32, d_model)
+            >>> memory_pos = torch.randn(20, 32, d_model)
+            >>> output = attention(curr, memory, curr_pos, memory_pos)
+            >>> print(output.shape)
+            torch.Size([10, 32, 256])
+        """
         if isinstance(curr, list):
             assert isinstance(curr_pos, list)
             assert len(curr) == len(curr_pos) == 1