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

ultralytics/nn/modules/block.py

@@ -1,6 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 """Block modules."""
 
+from typing import List, Optional, Tuple
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -60,15 +62,20 @@ class DFL(nn.Module):
     Proposed in Generalized Focal Loss https://ieeexplore.ieee.org/document/9792391
     """
 
-    def __init__(self, c1=16):
-        """Initialize a convolutional layer with a given number of input channels."""
+    def __init__(self, c1: int = 16):
+        """
+        Initialize a convolutional layer with a given number of input channels.
+
+        Args:
+            c1 (int): Number of input channels.
+        """
         super().__init__()
         self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
         x = torch.arange(c1, dtype=torch.float)
         self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1))
         self.c1 = c1
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Apply the DFL module to input tensor and return transformed output."""
         b, _, a = x.shape  # batch, channels, anchors
         return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)
@@ -78,7 +85,7 @@ class DFL(nn.Module):
 class Proto(nn.Module):
     """Ultralytics YOLO models mask Proto module for segmentation models."""
 
-    def __init__(self, c1, c_=256, c2=32):
+    def __init__(self, c1: int, c_: int = 256, c2: int = 32):
         """
         Initialize the Ultralytics YOLO models mask Proto module with specified number of protos and masks.
 
@@ -93,7 +100,7 @@ class Proto(nn.Module):
         self.cv2 = Conv(c_, c_, k=3)
         self.cv3 = Conv(c_, c2)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Perform a forward pass through layers using an upsampled input image."""
         return self.cv3(self.cv2(self.upsample(self.cv1(x))))
 
@@ -105,7 +112,7 @@ class HGStem(nn.Module):
     https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
     """
 
-    def __init__(self, c1, cm, c2):
+    def __init__(self, c1: int, cm: int, c2: int):
         """
         Initialize the StemBlock of PPHGNetV2.
 
@@ -122,7 +129,7 @@ class HGStem(nn.Module):
         self.stem4 = Conv(cm, c2, 1, 1, act=nn.ReLU())
         self.pool = nn.MaxPool2d(kernel_size=2, stride=1, padding=0, ceil_mode=True)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass of a PPHGNetV2 backbone layer."""
         x = self.stem1(x)
         x = F.pad(x, [0, 1, 0, 1])
@@ -143,7 +150,17 @@ class HGBlock(nn.Module):
     https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
     """
 
-    def __init__(self, c1, cm, c2, k=3, n=6, lightconv=False, shortcut=False, act=nn.ReLU()):
+    def __init__(
+        self,
+        c1: int,
+        cm: int,
+        c2: int,
+        k: int = 3,
+        n: int = 6,
+        lightconv: bool = False,
+        shortcut: bool = False,
+        act: nn.Module = nn.ReLU(),
+    ):
         """
         Initialize HGBlock with specified parameters.
 
@@ -164,7 +181,7 @@ class HGBlock(nn.Module):
         self.ec = Conv(c2 // 2, c2, 1, 1, act=act)  # excitation conv
         self.add = shortcut and c1 == c2
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass of a PPHGNetV2 backbone layer."""
         y = [x]
         y.extend(m(y[-1]) for m in self.m)
@@ -175,14 +192,14 @@ class HGBlock(nn.Module):
 class SPP(nn.Module):
     """Spatial Pyramid Pooling (SPP) layer https://arxiv.org/abs/1406.4729."""
 
-    def __init__(self, c1, c2, k=(5, 9, 13)):
+    def __init__(self, c1: int, c2: int, k: Tuple[int, ...] = (5, 9, 13)):
         """
         Initialize the SPP layer with input/output channels and pooling kernel sizes.
 
         Args:
             c1 (int): Input channels.
             c2 (int): Output channels.
-            k (Tuple[int, int, int]): Kernel sizes for max pooling.
+            k (tuple): Kernel sizes for max pooling.
         """
         super().__init__()
         c_ = c1 // 2  # hidden channels
@@ -190,7 +207,7 @@ class SPP(nn.Module):
         self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
         self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass of the SPP layer, performing spatial pyramid pooling."""
         x = self.cv1(x)
         return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
@@ -199,7 +216,7 @@ class SPP(nn.Module):
 class SPPF(nn.Module):
     """Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher."""
 
-    def __init__(self, c1, c2, k=5):
+    def __init__(self, c1: int, c2: int, k: int = 5):
         """
         Initialize the SPPF layer with given input/output channels and kernel size.
 
@@ -217,7 +234,7 @@ class SPPF(nn.Module):
         self.cv2 = Conv(c_ * 4, c2, 1, 1)
         self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Apply sequential pooling operations to input and return concatenated feature maps."""
         y = [self.cv1(x)]
         y.extend(self.m(y[-1]) for _ in range(3))
@@ -227,7 +244,7 @@ class SPPF(nn.Module):
 class C1(nn.Module):
     """CSP Bottleneck with 1 convolution."""
 
-    def __init__(self, c1, c2, n=1):
+    def __init__(self, c1: int, c2: int, n: int = 1):
         """
         Initialize the CSP Bottleneck with 1 convolution.
 
@@ -240,7 +257,7 @@ class C1(nn.Module):
         self.cv1 = Conv(c1, c2, 1, 1)
         self.m = nn.Sequential(*(Conv(c2, c2, 3) for _ in range(n)))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Apply convolution and residual connection to input tensor."""
         y = self.cv1(x)
         return self.m(y) + y
@@ -249,7 +266,7 @@ class C1(nn.Module):
 class C2(nn.Module):
     """CSP Bottleneck with 2 convolutions."""
 
-    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = True, g: int = 1, e: float = 0.5):
         """
         Initialize a CSP Bottleneck with 2 convolutions.
 
@@ -268,7 +285,7 @@ class C2(nn.Module):
         # self.attention = ChannelAttention(2 * self.c)  # or SpatialAttention()
         self.m = nn.Sequential(*(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n)))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass through the CSP bottleneck with 2 convolutions."""
         a, b = self.cv1(x).chunk(2, 1)
         return self.cv2(torch.cat((self.m(a), b), 1))
@@ -277,7 +294,7 @@ class C2(nn.Module):
 class C2f(nn.Module):
     """Faster Implementation of CSP Bottleneck with 2 convolutions."""
 
-    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
+    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = False, g: int = 1, e: float = 0.5):
         """
         Initialize a CSP bottleneck with 2 convolutions.
 
@@ -295,13 +312,13 @@ class C2f(nn.Module):
         self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
         self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass through C2f layer."""
         y = list(self.cv1(x).chunk(2, 1))
         y.extend(m(y[-1]) for m in self.m)
         return self.cv2(torch.cat(y, 1))
 
-    def forward_split(self, x):
+    def forward_split(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass using split() instead of chunk()."""
         y = self.cv1(x).split((self.c, self.c), 1)
         y = [y[0], y[1]]
@@ -312,7 +329,7 @@ class C2f(nn.Module):
 class C3(nn.Module):
     """CSP Bottleneck with 3 convolutions."""
 
-    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = True, g: int = 1, e: float = 0.5):
         """
         Initialize the CSP Bottleneck with 3 convolutions.
 
@@ -331,7 +348,7 @@ class C3(nn.Module):
         self.cv3 = Conv(2 * c_, c2, 1)  # optional act=FReLU(c2)
         self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, k=((1, 1), (3, 3)), e=1.0) for _ in range(n)))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass through the CSP bottleneck with 3 convolutions."""
         return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
 
@@ -339,7 +356,7 @@ class C3(nn.Module):
 class C3x(C3):
     """C3 module with cross-convolutions."""
 
-    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = True, g: int = 1, e: float = 0.5):
         """
         Initialize C3 module with cross-convolutions.
 
@@ -359,7 +376,7 @@ class C3x(C3):
 class RepC3(nn.Module):
     """Rep C3."""
 
-    def __init__(self, c1, c2, n=3, e=1.0):
+    def __init__(self, c1: int, c2: int, n: int = 3, e: float = 1.0):
         """
         Initialize CSP Bottleneck with a single convolution.
 
@@ -376,7 +393,7 @@ class RepC3(nn.Module):
         self.m = nn.Sequential(*[RepConv(c_, c_) for _ in range(n)])
         self.cv3 = Conv(c_, c2, 1, 1) if c_ != c2 else nn.Identity()
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass of RepC3 module."""
         return self.cv3(self.m(self.cv1(x)) + self.cv2(x))
 
@@ -384,7 +401,7 @@ class RepC3(nn.Module):
 class C3TR(C3):
     """C3 module with TransformerBlock()."""
 
-    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = True, g: int = 1, e: float = 0.5):
         """
         Initialize C3 module with TransformerBlock.
 
@@ -404,7 +421,7 @@ class C3TR(C3):
 class C3Ghost(C3):
     """C3 module with GhostBottleneck()."""
 
-    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = True, g: int = 1, e: float = 0.5):
         """
         Initialize C3 module with GhostBottleneck.
 
@@ -424,7 +441,7 @@ class C3Ghost(C3):
 class GhostBottleneck(nn.Module):
     """Ghost Bottleneck https://github.com/huawei-noah/Efficient-AI-Backbones."""
 
-    def __init__(self, c1, c2, k=3, s=1):
+    def __init__(self, c1: int, c2: int, k: int = 3, s: int = 1):
         """
         Initialize Ghost Bottleneck module.
 
@@ -445,7 +462,7 @@ class GhostBottleneck(nn.Module):
             nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Apply skip connection and concatenation to input tensor."""
         return self.conv(x) + self.shortcut(x)
 
@@ -453,7 +470,9 @@ class GhostBottleneck(nn.Module):
 class Bottleneck(nn.Module):
     """Standard bottleneck."""
 
-    def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
+    def __init__(
+        self, c1: int, c2: int, shortcut: bool = True, g: int = 1, k: Tuple[int, int] = (3, 3), e: float = 0.5
+    ):
         """
         Initialize a standard bottleneck module.
 
@@ -462,7 +481,7 @@ class Bottleneck(nn.Module):
             c2 (int): Output channels.
             shortcut (bool): Whether to use shortcut connection.
             g (int): Groups for convolutions.
-            k (Tuple[int, int]): Kernel sizes for convolutions.
+            k (tuple): Kernel sizes for convolutions.
             e (float): Expansion ratio.
         """
         super().__init__()
@@ -471,7 +490,7 @@ class Bottleneck(nn.Module):
         self.cv2 = Conv(c_, c2, k[1], 1, g=g)
         self.add = shortcut and c1 == c2
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Apply bottleneck with optional shortcut connection."""
         return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
 
@@ -479,7 +498,7 @@ class Bottleneck(nn.Module):
 class BottleneckCSP(nn.Module):
     """CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks."""
 
-    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = True, g: int = 1, e: float = 0.5):
         """
         Initialize CSP Bottleneck.
 
@@ -501,7 +520,7 @@ class BottleneckCSP(nn.Module):
         self.act = nn.SiLU()
         self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Apply CSP bottleneck with 3 convolutions."""
         y1 = self.cv3(self.m(self.cv1(x)))
         y2 = self.cv2(x)
@@ -511,7 +530,7 @@ class BottleneckCSP(nn.Module):
 class ResNetBlock(nn.Module):
     """ResNet block with standard convolution layers."""
 
-    def __init__(self, c1, c2, s=1, e=4):
+    def __init__(self, c1: int, c2: int, s: int = 1, e: int = 4):
         """
         Initialize ResNet block.
 
@@ -528,7 +547,7 @@ class ResNetBlock(nn.Module):
         self.cv3 = Conv(c2, c3, k=1, act=False)
         self.shortcut = nn.Sequential(Conv(c1, c3, k=1, s=s, act=False)) if s != 1 or c1 != c3 else nn.Identity()
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass through the ResNet block."""
         return F.relu(self.cv3(self.cv2(self.cv1(x))) + self.shortcut(x))
 
@@ -536,7 +555,7 @@ class ResNetBlock(nn.Module):
 class ResNetLayer(nn.Module):
     """ResNet layer with multiple ResNet blocks."""
 
-    def __init__(self, c1, c2, s=1, is_first=False, n=1, e=4):
+    def __init__(self, c1: int, c2: int, s: int = 1, is_first: bool = False, n: int = 1, e: int = 4):
         """
         Initialize ResNet layer.
 
@@ -560,7 +579,7 @@ class ResNetLayer(nn.Module):
             blocks.extend([ResNetBlock(e * c2, c2, 1, e=e) for _ in range(n - 1)])
             self.layer = nn.Sequential(*blocks)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass through the ResNet layer."""
         return self.layer(x)
 
@@ -568,7 +587,7 @@ class ResNetLayer(nn.Module):
 class MaxSigmoidAttnBlock(nn.Module):
     """Max Sigmoid attention block."""
 
-    def __init__(self, c1, c2, nh=1, ec=128, gc=512, scale=False):
+    def __init__(self, c1: int, c2: int, nh: int = 1, ec: int = 128, gc: int = 512, scale: bool = False):
         """
         Initialize MaxSigmoidAttnBlock.
 
@@ -589,7 +608,7 @@ class MaxSigmoidAttnBlock(nn.Module):
         self.proj_conv = Conv(c1, c2, k=3, s=1, act=False)
         self.scale = nn.Parameter(torch.ones(1, nh, 1, 1)) if scale else 1.0
 
-    def forward(self, x, guide):
+    def forward(self, x: torch.Tensor, guide: torch.Tensor) -> torch.Tensor:
         """
         Forward pass of MaxSigmoidAttnBlock.
 
@@ -622,7 +641,18 @@ class MaxSigmoidAttnBlock(nn.Module):
 class C2fAttn(nn.Module):
     """C2f module with an additional attn module."""
 
-    def __init__(self, c1, c2, n=1, ec=128, nh=1, gc=512, shortcut=False, g=1, e=0.5):
+    def __init__(
+        self,
+        c1: int,
+        c2: int,
+        n: int = 1,
+        ec: int = 128,
+        nh: int = 1,
+        gc: int = 512,
+        shortcut: bool = False,
+        g: int = 1,
+        e: float = 0.5,
+    ):
         """
         Initialize C2f module with attention mechanism.
 
@@ -644,7 +674,7 @@ class C2fAttn(nn.Module):
         self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
         self.attn = MaxSigmoidAttnBlock(self.c, self.c, gc=gc, ec=ec, nh=nh)
 
-    def forward(self, x, guide):
+    def forward(self, x: torch.Tensor, guide: torch.Tensor) -> torch.Tensor:
         """
         Forward pass through C2f layer with attention.
 
@@ -660,7 +690,7 @@ class C2fAttn(nn.Module):
         y.append(self.attn(y[-1], guide))
         return self.cv2(torch.cat(y, 1))
 
-    def forward_split(self, x, guide):
+    def forward_split(self, x: torch.Tensor, guide: torch.Tensor) -> torch.Tensor:
         """
         Forward pass using split() instead of chunk().
 
@@ -680,7 +710,9 @@ class C2fAttn(nn.Module):
 class ImagePoolingAttn(nn.Module):
     """ImagePoolingAttn: Enhance the text embeddings with image-aware information."""
 
-    def __init__(self, ec=256, ch=(), ct=512, nh=8, k=3, scale=False):
+    def __init__(
+        self, ec: int = 256, ch: Tuple[int, ...] = (), ct: int = 512, nh: int = 8, k: int = 3, scale: bool = False
+    ):
         """
         Initialize ImagePoolingAttn module.
 
@@ -708,7 +740,7 @@ class ImagePoolingAttn(nn.Module):
         self.hc = ec // nh
         self.k = k
 
-    def forward(self, x, text):
+    def forward(self, x: List[torch.Tensor], text: torch.Tensor) -> torch.Tensor:
         """
         Forward pass of ImagePoolingAttn.
 
@@ -752,7 +784,7 @@ class ContrastiveHead(nn.Module):
         self.bias = nn.Parameter(torch.tensor([-10.0]))
         self.logit_scale = nn.Parameter(torch.ones([]) * torch.tensor(1 / 0.07).log())
 
-    def forward(self, x, w):
+    def forward(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
         """
         Forward function of contrastive learning.
 
@@ -798,15 +830,11 @@ class BNContrastiveHead(nn.Module):
         del self.logit_scale
         self.forward = self.forward_fuse
 
-    def forward_fuse(self, x, w):
-        """
-        Passes input out unchanged.
-
-        TODO: Update or remove?
-        """
+    def forward_fuse(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
+        """Passes input out unchanged."""
         return x
 
-    def forward(self, x, w):
+    def forward(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
         """
         Forward function of contrastive learning with batch normalization.
 
@@ -827,7 +855,9 @@ class BNContrastiveHead(nn.Module):
 class RepBottleneck(Bottleneck):
     """Rep bottleneck."""
 
-    def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
+    def __init__(
+        self, c1: int, c2: int, shortcut: bool = True, g: int = 1, k: Tuple[int, int] = (3, 3), e: float = 0.5
+    ):
         """
         Initialize RepBottleneck.
 
@@ -836,7 +866,7 @@ class RepBottleneck(Bottleneck):
             c2 (int): Output channels.
             shortcut (bool): Whether to use shortcut connection.
             g (int): Groups for convolutions.
-            k (Tuple[int, int]): Kernel sizes for convolutions.
+            k (tuple): Kernel sizes for convolutions.
             e (float): Expansion ratio.
         """
         super().__init__(c1, c2, shortcut, g, k, e)
@@ -847,7 +877,7 @@ class RepBottleneck(Bottleneck):
 class RepCSP(C3):
     """Repeatable Cross Stage Partial Network (RepCSP) module for efficient feature extraction."""
 
-    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = True, g: int = 1, e: float = 0.5):
         """
         Initialize RepCSP layer.
 
@@ -867,7 +897,7 @@ class RepCSP(C3):
 class RepNCSPELAN4(nn.Module):
     """CSP-ELAN."""
 
-    def __init__(self, c1, c2, c3, c4, n=1):
+    def __init__(self, c1: int, c2: int, c3: int, c4: int, n: int = 1):
         """
         Initialize CSP-ELAN layer.
 
@@ -885,13 +915,13 @@ class RepNCSPELAN4(nn.Module):
         self.cv3 = nn.Sequential(RepCSP(c4, c4, n), Conv(c4, c4, 3, 1))
         self.cv4 = Conv(c3 + (2 * c4), c2, 1, 1)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass through RepNCSPELAN4 layer."""
         y = list(self.cv1(x).chunk(2, 1))
         y.extend((m(y[-1])) for m in [self.cv2, self.cv3])
         return self.cv4(torch.cat(y, 1))
 
-    def forward_split(self, x):
+    def forward_split(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass using split() instead of chunk()."""
         y = list(self.cv1(x).split((self.c, self.c), 1))
         y.extend(m(y[-1]) for m in [self.cv2, self.cv3])
@@ -901,7 +931,7 @@ class RepNCSPELAN4(nn.Module):
 class ELAN1(RepNCSPELAN4):
     """ELAN1 module with 4 convolutions."""
 
-    def __init__(self, c1, c2, c3, c4):
+    def __init__(self, c1: int, c2: int, c3: int, c4: int):
         """
         Initialize ELAN1 layer.
 
@@ -922,7 +952,7 @@ class ELAN1(RepNCSPELAN4):
 class AConv(nn.Module):
     """AConv."""
 
-    def __init__(self, c1, c2):
+    def __init__(self, c1: int, c2: int):
         """
         Initialize AConv module.
 
@@ -933,7 +963,7 @@ class AConv(nn.Module):
         super().__init__()
         self.cv1 = Conv(c1, c2, 3, 2, 1)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass through AConv layer."""
         x = torch.nn.functional.avg_pool2d(x, 2, 1, 0, False, True)
         return self.cv1(x)
@@ -942,7 +972,7 @@ class AConv(nn.Module):
 class ADown(nn.Module):
     """ADown."""
 
-    def __init__(self, c1, c2):
+    def __init__(self, c1: int, c2: int):
         """
         Initialize ADown module.
 
@@ -955,7 +985,7 @@ class ADown(nn.Module):
         self.cv1 = Conv(c1 // 2, self.c, 3, 2, 1)
         self.cv2 = Conv(c1 // 2, self.c, 1, 1, 0)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass through ADown layer."""
         x = torch.nn.functional.avg_pool2d(x, 2, 1, 0, False, True)
         x1, x2 = x.chunk(2, 1)
@@ -968,7 +998,7 @@ class ADown(nn.Module):
 class SPPELAN(nn.Module):
     """SPP-ELAN."""
 
-    def __init__(self, c1, c2, c3, k=5):
+    def __init__(self, c1: int, c2: int, c3: int, k: int = 5):
         """
         Initialize SPP-ELAN block.
 
@@ -986,7 +1016,7 @@ class SPPELAN(nn.Module):
         self.cv4 = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
         self.cv5 = Conv(4 * c3, c2, 1, 1)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass through SPPELAN layer."""
         y = [self.cv1(x)]
         y.extend(m(y[-1]) for m in [self.cv2, self.cv3, self.cv4])
@@ -996,7 +1026,7 @@ class SPPELAN(nn.Module):
 class CBLinear(nn.Module):
     """CBLinear."""
 
-    def __init__(self, c1, c2s, k=1, s=1, p=None, g=1):
+    def __init__(self, c1: int, c2s: List[int], k: int = 1, s: int = 1, p: Optional[int] = None, g: int = 1):
         """
         Initialize CBLinear module.
 
@@ -1012,7 +1042,7 @@ class CBLinear(nn.Module):
         self.c2s = c2s
         self.conv = nn.Conv2d(c1, sum(c2s), k, s, autopad(k, p), groups=g, bias=True)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
         """Forward pass through CBLinear layer."""
         return self.conv(x).split(self.c2s, dim=1)
 
@@ -1020,7 +1050,7 @@ class CBLinear(nn.Module):
 class CBFuse(nn.Module):
     """CBFuse."""
 
-    def __init__(self, idx):
+    def __init__(self, idx: List[int]):
         """
         Initialize CBFuse module.
 
@@ -1030,7 +1060,7 @@ class CBFuse(nn.Module):
         super().__init__()
         self.idx = idx
 
-    def forward(self, xs):
+    def forward(self, xs: List[torch.Tensor]) -> torch.Tensor:
         """
         Forward pass through CBFuse layer.
 
@@ -1048,7 +1078,7 @@ class CBFuse(nn.Module):
 class C3f(nn.Module):
     """Faster Implementation of CSP Bottleneck with 2 convolutions."""
 
-    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
+    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = False, g: int = 1, e: float = 0.5):
         """
         Initialize CSP bottleneck layer with two convolutions.
 
@@ -1067,7 +1097,7 @@ class C3f(nn.Module):
         self.cv3 = Conv((2 + n) * c_, c2, 1)  # optional act=FReLU(c2)
         self.m = nn.ModuleList(Bottleneck(c_, c_, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass through C3f layer."""
         y = [self.cv2(x), self.cv1(x)]
         y.extend(m(y[-1]) for m in self.m)
@@ -1077,7 +1107,9 @@ class C3f(nn.Module):
 class C3k2(C2f):
     """Faster Implementation of CSP Bottleneck with 2 convolutions."""
 
-    def __init__(self, c1, c2, n=1, c3k=False, e=0.5, g=1, shortcut=True):
+    def __init__(
+        self, c1: int, c2: int, n: int = 1, c3k: bool = False, e: float = 0.5, g: int = 1, shortcut: bool = True
+    ):
         """
         Initialize C3k2 module.
 
@@ -1099,7 +1131,7 @@ class C3k2(C2f):
 class C3k(C3):
     """C3k is a CSP bottleneck module with customizable kernel sizes for feature extraction in neural networks."""
 
-    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, k=3):
+    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = True, g: int = 1, e: float = 0.5, k: int = 3):
         """
         Initialize C3k module.
 
@@ -1121,7 +1153,7 @@ class C3k(C3):
 class RepVGGDW(torch.nn.Module):
     """RepVGGDW is a class that represents a depth wise separable convolutional block in RepVGG architecture."""
 
-    def __init__(self, ed) -> None:
+    def __init__(self, ed: int) -> None:
         """
         Initialize RepVGGDW module.
 
@@ -1134,7 +1166,7 @@ class RepVGGDW(torch.nn.Module):
         self.dim = ed
         self.act = nn.SiLU()
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Perform a forward pass of the RepVGGDW block.
 
@@ -1146,7 +1178,7 @@ class RepVGGDW(torch.nn.Module):
         """
         return self.act(self.conv(x) + self.conv1(x))
 
-    def forward_fuse(self, x):
+    def forward_fuse(self, x: torch.Tensor) -> torch.Tensor:
         """
         Perform a forward pass of the RepVGGDW block without fusing the convolutions.
 
@@ -1197,7 +1229,7 @@ class CIB(nn.Module):
         lk (bool, optional): Whether to use RepVGGDW for the third convolutional layer. Defaults to False.
     """
 
-    def __init__(self, c1, c2, shortcut=True, e=0.5, lk=False):
+    def __init__(self, c1: int, c2: int, shortcut: bool = True, e: float = 0.5, lk: bool = False):
         """
         Initialize the CIB module.
 
@@ -1220,7 +1252,7 @@ class CIB(nn.Module):
 
         self.add = shortcut and c1 == c2
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Forward pass of the CIB module.
 
@@ -1247,7 +1279,9 @@ class C2fCIB(C2f):
         e (float, optional): Expansion ratio for CIB modules. Defaults to 0.5.
     """
 
-    def __init__(self, c1, c2, n=1, shortcut=False, lk=False, g=1, e=0.5):
+    def __init__(
+        self, c1: int, c2: int, n: int = 1, shortcut: bool = False, lk: bool = False, g: int = 1, e: float = 0.5
+    ):
         """
         Initialize C2fCIB module.
 
@@ -1283,7 +1317,7 @@ class Attention(nn.Module):
         pe (Conv): Convolutional layer for positional encoding.
     """
 
-    def __init__(self, dim, num_heads=8, attn_ratio=0.5):
+    def __init__(self, dim: int, num_heads: int = 8, attn_ratio: float = 0.5):
         """
         Initialize multi-head attention module.
 
@@ -1303,7 +1337,7 @@ class Attention(nn.Module):
         self.proj = Conv(dim, dim, 1, act=False)
         self.pe = Conv(dim, dim, 3, 1, g=dim, act=False)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Forward pass of the Attention module.
 
@@ -1349,7 +1383,7 @@ class PSABlock(nn.Module):
         >>> output_tensor = psablock(input_tensor)
     """
 
-    def __init__(self, c, attn_ratio=0.5, num_heads=4, shortcut=True) -> None:
+    def __init__(self, c: int, attn_ratio: float = 0.5, num_heads: int = 4, shortcut: bool = True) -> None:
         """
         Initialize the PSABlock.
 
@@ -1365,7 +1399,7 @@ class PSABlock(nn.Module):
         self.ffn = nn.Sequential(Conv(c, c * 2, 1), Conv(c * 2, c, 1, act=False))
         self.add = shortcut
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Execute a forward pass through PSABlock.
 
@@ -1404,7 +1438,7 @@ class PSA(nn.Module):
         >>> output_tensor = psa.forward(input_tensor)
     """
 
-    def __init__(self, c1, c2, e=0.5):
+    def __init__(self, c1: int, c2: int, e: float = 0.5):
         """
         Initialize PSA module.
 
@@ -1422,7 +1456,7 @@ class PSA(nn.Module):
         self.attn = Attention(self.c, attn_ratio=0.5, num_heads=self.c // 64)
         self.ffn = nn.Sequential(Conv(self.c, self.c * 2, 1), Conv(self.c * 2, self.c, 1, act=False))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Execute forward pass in PSA module.
 
@@ -1463,7 +1497,7 @@ class C2PSA(nn.Module):
         >>> output_tensor = c2psa(input_tensor)
     """
 
-    def __init__(self, c1, c2, n=1, e=0.5):
+    def __init__(self, c1: int, c2: int, n: int = 1, e: float = 0.5):
         """
         Initialize C2PSA module.
 
@@ -1481,7 +1515,7 @@ class C2PSA(nn.Module):
 
         self.m = nn.Sequential(*(PSABlock(self.c, attn_ratio=0.5, num_heads=self.c // 64) for _ in range(n)))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Process the input tensor through a series of PSA blocks.
 
@@ -1521,7 +1555,7 @@ class C2fPSA(C2f):
         >>> print(output.shape)
     """
 
-    def __init__(self, c1, c2, n=1, e=0.5):
+    def __init__(self, c1: int, c2: int, n: int = 1, e: float = 0.5):
         """
         Initialize C2fPSA module.
 
@@ -1560,7 +1594,7 @@ class SCDown(nn.Module):
         torch.Size([1, 128, 64, 64])
     """
 
-    def __init__(self, c1, c2, k, s):
+    def __init__(self, c1: int, c2: int, k: int, s: int):
         """
         Initialize SCDown module.
 
@@ -1574,7 +1608,7 @@ class SCDown(nn.Module):
         self.cv1 = Conv(c1, c2, 1, 1)
         self.cv2 = Conv(c2, c2, k=k, s=s, g=c2, act=False)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Apply convolution and downsampling to the input tensor.
 
@@ -1604,7 +1638,9 @@ class TorchVision(nn.Module):
         split (bool, optional): Returns output from intermediate child modules as list. Default is False.
     """
 
-    def __init__(self, model, weights="DEFAULT", unwrap=True, truncate=2, split=False):
+    def __init__(
+        self, model: str, weights: str = "DEFAULT", unwrap: bool = True, truncate: int = 2, split: bool = False
+    ):
         """
         Load the model and weights from torchvision.
 
@@ -1632,7 +1668,7 @@ class TorchVision(nn.Module):
             self.split = False
             self.m.head = self.m.heads = nn.Identity()
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Forward pass through the model.
 
@@ -1676,14 +1712,14 @@ class AAttn(nn.Module):
         torch.Size([1, 256, 32, 32])
     """
 
-    def __init__(self, dim, num_heads, area=1):
+    def __init__(self, dim: int, num_heads: int, area: int = 1):
         """
         Initialize an Area-attention module for YOLO models.
 
         Args:
             dim (int): Number of hidden channels.
             num_heads (int): Number of heads into which the attention mechanism is divided.
-            area (int): Number of areas the feature map is divided, default is 1.
+            area (int): Number of areas the feature map is divided.
         """
         super().__init__()
         self.area = area
@@ -1696,7 +1732,7 @@ class AAttn(nn.Module):
         self.proj = Conv(all_head_dim, dim, 1, act=False)
         self.pe = Conv(all_head_dim, dim, 7, 1, 3, g=dim, act=False)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Process the input tensor through the area-attention.
 
@@ -1760,7 +1796,7 @@ class ABlock(nn.Module):
         torch.Size([1, 256, 32, 32])
     """
 
-    def __init__(self, dim, num_heads, mlp_ratio=1.2, area=1):
+    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 1.2, area: int = 1):
         """
         Initialize an Area-attention block module.
 
@@ -1778,7 +1814,7 @@ class ABlock(nn.Module):
 
         self.apply(self._init_weights)
 
-    def _init_weights(self, m):
+    def _init_weights(self, m: nn.Module):
         """
         Initialize weights using a truncated normal distribution.
 
@@ -1790,7 +1826,7 @@ class ABlock(nn.Module):
             if m.bias is not None:
                 nn.init.constant_(m.bias, 0)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Forward pass through ABlock.
 
@@ -1828,7 +1864,19 @@ class A2C2f(nn.Module):
         torch.Size([1, 512, 32, 32])
     """
 
-    def __init__(self, c1, c2, n=1, a2=True, area=1, residual=False, mlp_ratio=2.0, e=0.5, g=1, shortcut=True):
+    def __init__(
+        self,
+        c1: int,
+        c2: int,
+        n: int = 1,
+        a2: bool = True,
+        area: int = 1,
+        residual: bool = False,
+        mlp_ratio: float = 2.0,
+        e: float = 0.5,
+        g: int = 1,
+        shortcut: bool = True,
+    ):
         """
         Initialize Area-Attention C2f module.
 
@@ -1859,7 +1907,7 @@ class A2C2f(nn.Module):
             for _ in range(n)
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Forward pass through A2C2f layer.
 
@@ -1880,13 +1928,20 @@ class A2C2f(nn.Module):
 class SwiGLUFFN(nn.Module):
     """SwiGLU Feed-Forward Network for transformer-based architectures."""
 
-    def __init__(self, gc, ec, e=4) -> None:
-        """Initialize SwiGLU FFN with input dimension, output dimension, and expansion factor."""
+    def __init__(self, gc: int, ec: int, e: int = 4) -> None:
+        """
+        Initialize SwiGLU FFN with input dimension, output dimension, and expansion factor.
+
+        Args:
+            gc (int): Guide channels.
+            ec (int): Embedding channels.
+            e (int): Expansion factor.
+        """
         super().__init__()
         self.w12 = nn.Linear(gc, e * ec)
         self.w3 = nn.Linear(e * ec // 2, ec)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Apply SwiGLU transformation to input features."""
         x12 = self.w12(x)
         x1, x2 = x12.chunk(2, dim=-1)
@@ -1897,8 +1952,13 @@ class SwiGLUFFN(nn.Module):
 class Residual(nn.Module):
     """Residual connection wrapper for neural network modules."""
 
-    def __init__(self, m) -> None:
-        """Initialize residual module with the wrapped module."""
+    def __init__(self, m: nn.Module) -> None:
+        """
+        Initialize residual module with the wrapped module.
+
+        Args:
+            m (nn.Module): Module to wrap with residual connection.
+        """
         super().__init__()
         self.m = m
         nn.init.zeros_(self.m.w3.bias)
@@ -1906,7 +1966,7 @@ class Residual(nn.Module):
         # nn.init.constant_(self.m.w3.weight, 1e-6)
         nn.init.zeros_(self.m.w3.weight)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Apply residual connection to input features."""
         return x + self.m(x)
 
@@ -1914,8 +1974,15 @@ class Residual(nn.Module):
 class SAVPE(nn.Module):
     """Spatial-Aware Visual Prompt Embedding module for feature enhancement."""
 
-    def __init__(self, ch, c3, embed):
-        """Initialize SAVPE module with channels, intermediate channels, and embedding dimension."""
+    def __init__(self, ch: List[int], c3: int, embed: int):
+        """
+        Initialize SAVPE module with channels, intermediate channels, and embedding dimension.
+
+        Args:
+            ch (List[int]): List of input channel dimensions.
+            c3 (int): Intermediate channels.
+            embed (int): Embedding dimension.
+        """
         super().__init__()
         self.cv1 = nn.ModuleList(
             nn.Sequential(
@@ -1935,7 +2002,7 @@ class SAVPE(nn.Module):
         self.cv5 = nn.Conv2d(1, self.c, 3, padding=1)
         self.cv6 = nn.Sequential(Conv(2 * self.c, self.c, 3), nn.Conv2d(self.c, self.c, 3, padding=1))
 
-    def forward(self, x, vp):
+    def forward(self, x: List[torch.Tensor], vp: torch.Tensor) -> torch.Tensor:
         """Process input features and visual prompts to generate enhanced embeddings."""
         y = [self.cv2[i](xi) for i, xi in enumerate(x)]
         y = self.cv4(torch.cat(y, dim=1))