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

ultralytics/models/sam/modules/transformer.py

@@ -27,7 +27,7 @@ class TwoWayTransformer(nn.Module):
         norm_final_attn (nn.LayerNorm): Layer normalization applied to final queries.
 
     Methods:
-        forward: Processes image and point embeddings through the transformer.
+        forward: Process image and point embeddings through the transformer.
 
     Examples:
         >>> transformer = TwoWayTransformer(depth=6, embedding_dim=256, num_heads=8, mlp_dim=2048)
@@ -55,8 +55,8 @@ class TwoWayTransformer(nn.Module):
             embedding_dim (int): Channel dimension for input embeddings.
             num_heads (int): Number of heads for multihead attention. Must divide embedding_dim.
             mlp_dim (int): Internal channel dimension for the MLP block.
-            activation (Type[nn.Module]): Activation function to use in the MLP block.
-            attention_downsample_rate (int): Downsampling rate for attention mechanism.
+            activation (Type[nn.Module], optional): Activation function to use in the MLP block.
+            attention_downsample_rate (int, optional): Downsampling rate for attention mechanism.
         """
         super().__init__()
         self.depth = depth
@@ -145,7 +145,7 @@ class TwoWayAttentionBlock(nn.Module):
         skip_first_layer_pe (bool): Whether to skip positional encoding in the first layer.
 
     Methods:
-        forward: Applies self-attention and cross-attention to queries and keys.
+        forward: Apply self-attention and cross-attention to queries and keys.
 
     Examples:
         >>> embedding_dim, num_heads = 256, 8
@@ -176,10 +176,10 @@ class TwoWayAttentionBlock(nn.Module):
         Args:
             embedding_dim (int): Channel dimension of the embeddings.
             num_heads (int): Number of attention heads in the attention layers.
-            mlp_dim (int): Hidden dimension of the MLP block.
-            activation (Type[nn.Module]): Activation function for the MLP block.
-            attention_downsample_rate (int): Downsampling rate for the attention mechanism.
-            skip_first_layer_pe (bool): Whether to skip positional encoding in the first layer.
+            mlp_dim (int, optional): Hidden dimension of the MLP block.
+            activation (Type[nn.Module], optional): Activation function for the MLP block.
+            attention_downsample_rate (int, optional): Downsampling rate for the attention mechanism.
+            skip_first_layer_pe (bool, optional): Whether to skip positional encoding in the first layer.
         """
         super().__init__()
         self.self_attn = Attention(embedding_dim, num_heads)
@@ -259,9 +259,9 @@ class Attention(nn.Module):
         out_proj (nn.Linear): Linear projection for output.
 
     Methods:
-        _separate_heads: Separates input tensor into attention heads.
-        _recombine_heads: Recombines separated attention heads.
-        forward: Computes attention output for given query, key, and value tensors.
+        _separate_heads: Separate input tensor into attention heads.
+        _recombine_heads: Recombine separated attention heads.
+        forward: Compute attention output for given query, key, and value tensors.
 
     Examples:
         >>> attn = Attention(embedding_dim=256, num_heads=8, downsample_rate=2)
@@ -285,8 +285,8 @@ class Attention(nn.Module):
         Args:
             embedding_dim (int): Dimensionality of input embeddings.
             num_heads (int): Number of attention heads.
-            downsample_rate (int): Factor by which internal dimensions are downsampled.
-            kv_in_dim (int | None): Dimensionality of key and value inputs. If None, uses embedding_dim.
+            downsample_rate (int, optional): Factor by which internal dimensions are downsampled.
+            kv_in_dim (int | None, optional): Dimensionality of key and value inputs. If None, uses embedding_dim.
 
         Raises:
             AssertionError: If num_heads does not evenly divide the internal dim (embedding_dim / downsample_rate).

ultralytics/models/sam/modules/utils.py

@@ -1,12 +1,12 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from typing import Tuple
+from typing import Any, Dict, Tuple
 
 import torch
 import torch.nn.functional as F
 
 
-def select_closest_cond_frames(frame_idx, cond_frame_outputs, max_cond_frame_num):
+def select_closest_cond_frames(frame_idx: int, cond_frame_outputs: Dict[int, Any], max_cond_frame_num: int):
     """
     Select the closest conditioning frames to a given frame index.
 
@@ -59,14 +59,14 @@ def select_closest_cond_frames(frame_idx, cond_frame_outputs, max_cond_frame_num
     return selected_outputs, unselected_outputs
 
 
-def get_1d_sine_pe(pos_inds, dim, temperature=10000):
+def get_1d_sine_pe(pos_inds: torch.Tensor, dim: int, temperature: float = 10000):
     """
     Generate 1D sinusoidal positional embeddings for given positions and dimensions.
 
     Args:
         pos_inds (torch.Tensor): Position indices for which to generate embeddings.
         dim (int): Dimension of the positional embeddings. Should be an even number.
-        temperature (float): Scaling factor for the frequency of the sinusoidal functions.
+        temperature (float, optional): Scaling factor for the frequency of the sinusoidal functions.
 
     Returns:
         (torch.Tensor): Sinusoidal positional embeddings with shape (pos_inds.shape, dim).
@@ -98,14 +98,11 @@ def init_t_xy(end_x: int, end_y: int):
         end_y (int): Height of the grid (number of rows).
 
     Returns:
-        t (torch.Tensor): Linear indices for each position in the grid, with shape (end_x * end_y).
         t_x (torch.Tensor): X-coordinates for each position, with shape (end_x * end_y).
         t_y (torch.Tensor): Y-coordinates for each position, with shape (end_x * end_y).
 
     Examples:
-        >>> t, t_x, t_y = init_t_xy(3, 2)
-        >>> print(t)
-        tensor([0., 1., 2., 3., 4., 5.])
+        >>> t_x, t_y = init_t_xy(3, 2)
         >>> print(t_x)
         tensor([0., 1., 2., 0., 1., 2.])
         >>> print(t_y)
@@ -131,18 +128,13 @@ def compute_axial_cis(dim: int, end_x: int, end_y: int, theta: float = 10000.0):
         theta (float, optional): Scaling factor for frequency computation.
 
     Returns:
-        freqs_cis_x (torch.Tensor): Complex exponential positional encodings for x-dimension with shape
-            (end_x*end_y, dim//4).
-        freqs_cis_y (torch.Tensor): Complex exponential positional encodings for y-dimension with shape
-            (end_x*end_y, dim//4).
+        (torch.Tensor): Complex exponential positional encodings with shape (end_x*end_y, dim//2).
 
     Examples:
         >>> dim, end_x, end_y = 128, 8, 8
-        >>> freqs_cis_x, freqs_cis_y = compute_axial_cis(dim, end_x, end_y)
-        >>> freqs_cis_x.shape
-        torch.Size([64, 32])
-        >>> freqs_cis_y.shape
-        torch.Size([64, 32])
+        >>> freqs_cis = compute_axial_cis(dim, end_x, end_y)
+        >>> freqs_cis.shape
+        torch.Size([64, 64])
     """
     freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
     freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
@@ -225,7 +217,7 @@ def apply_rotary_enc(
     return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)
 
 
-def window_partition(x, window_size):
+def window_partition(x: torch.Tensor, window_size: int):
     """
     Partition input tensor into non-overlapping windows with padding if needed.
 
@@ -256,7 +248,7 @@ def window_partition(x, window_size):
     return windows, (Hp, Wp)
 
 
-def window_unpartition(windows, window_size, pad_hw, hw):
+def window_unpartition(windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int]):
     """
     Unpartition windowed sequences into original sequences and remove padding.