wavedl 1.6.0__py3-none-any.whl → 1.6.1__py3-none-any.whl

wavedl/models/__init__.py

@@ -80,8 +80,24 @@ from .vit import ViTBase_, ViTSmall, ViTTiny
 # Optional timm-based models (imported conditionally)
 try:
     from .caformer import CaFormerS18, CaFormerS36, PoolFormerS12
+    from .efficientvit import (
+        EfficientViTB0,
+        EfficientViTB1,
+        EfficientViTB2,
+        EfficientViTB3,
+        EfficientViTL1,
+        EfficientViTL2,
+        EfficientViTM0,
+        EfficientViTM1,
+        EfficientViTM2,
+    )
     from .fastvit import FastViTS12, FastViTSA12, FastViTT8, FastViTT12
     from .maxvit import MaxViTBaseLarge, MaxViTSmall, MaxViTTiny
+    from .unireplknet import (
+        UniRepLKNetBaseLarge,
+        UniRepLKNetSmall,
+        UniRepLKNetTiny,
+    )
 
     _HAS_TIMM_MODELS = True
 except ImportError:
@@ -148,6 +164,15 @@ if _HAS_TIMM_MODELS:
         [
             "CaFormerS18",
             "CaFormerS36",
+            "EfficientViTB0",
+            "EfficientViTB1",
+            "EfficientViTB2",
+            "EfficientViTB3",
+            "EfficientViTL1",
+            "EfficientViTL2",
+            "EfficientViTM0",
+            "EfficientViTM1",
+            "EfficientViTM2",
             "FastViTS12",
             "FastViTSA12",
             "FastViTT8",
@@ -156,5 +181,8 @@ if _HAS_TIMM_MODELS:
             "MaxViTSmall",
             "MaxViTTiny",
             "PoolFormerS12",
+            "UniRepLKNetBaseLarge",
+            "UniRepLKNetSmall",
+            "UniRepLKNetTiny",
         ]
     )

wavedl/models/{_timm_utils.py → _pretrained_utils.py}

@@ -236,3 +236,131 @@ def adapt_input_channels(
         return new_conv
     else:
         raise NotImplementedError(f"Unsupported layer type: {type(conv_layer)}")
+
+
+def adapt_first_conv_for_single_channel(
+    module: nn.Module,
+    conv_path: str,
+    pretrained: bool = True,
+) -> None:
+    """
+    Adapt the first convolutional layer of a pretrained model for single-channel input.
+
+    This is a convenience function for torchvision-style models where the path
+    to the first conv layer is known. It modifies the model in-place.
+
+    For pretrained models, the RGB weights are averaged to create grayscale weights,
+    which provides a reasonable initialization for single-channel inputs.
+
+    Args:
+        module: The model or submodule containing the conv layer
+        conv_path: Dot-separated path to the conv layer (e.g., "conv1", "features.0.0")
+        pretrained: Whether to adapt pretrained weights by averaging RGB channels
+
+    Example:
+        >>> # For torchvision ResNet
+        >>> adapt_first_conv_for_single_channel(
+        ...     model.backbone, "conv1", pretrained=True
+        ... )
+        >>> # For torchvision ConvNeXt
+        >>> adapt_first_conv_for_single_channel(
+        ...     model.backbone, "features.0.0", pretrained=True
+        ... )
+        >>> # For torchvision DenseNet
+        >>> adapt_first_conv_for_single_channel(
+        ...     model.backbone, "features.conv0", pretrained=True
+        ... )
+    """
+    # Navigate to parent and get the conv layer
+    parts = conv_path.split(".")
+    parent = module
+    for part in parts[:-1]:
+        if part.isdigit():
+            parent = parent[int(part)]
+        else:
+            parent = getattr(parent, part)
+
+    # Get the final attribute name and the old conv
+    final_attr = parts[-1]
+    if final_attr.isdigit():
+        old_conv = parent[int(final_attr)]
+    else:
+        old_conv = getattr(parent, final_attr)
+
+    # Create and set the new conv
+    new_conv = adapt_input_channels(old_conv, new_in_channels=1, pretrained=pretrained)
+
+    if final_attr.isdigit():
+        parent[int(final_attr)] = new_conv
+    else:
+        setattr(parent, final_attr, new_conv)
+
+
+def find_and_adapt_input_convs(
+    backbone: nn.Module,
+    pretrained: bool = True,
+    adapt_all: bool = False,
+) -> int:
+    """
+    Find and adapt Conv2d layers with 3 input channels for single-channel input.
+
+    This is useful for timm-style models where the exact path to the first
+    conv layer may vary or where multiple layers need adaptation.
+
+    Args:
+        backbone: The backbone model to adapt
+        pretrained: Whether to adapt pretrained weights by averaging RGB channels
+        adapt_all: If True, adapt all Conv2d layers with 3 input channels.
+                   If False (default), only adapt the first one found.
+
+    Returns:
+        Number of layers adapted
+
+    Example:
+        >>> # For timm models (adapt first conv only)
+        >>> count = find_and_adapt_input_convs(model.backbone, pretrained=True)
+        >>> # For models with multiple input convs (e.g., FastViT)
+        >>> count = find_and_adapt_input_convs(
+        ...     model.backbone, pretrained=True, adapt_all=True
+        ... )
+    """
+    adapted_count = 0
+
+    for name, module in backbone.named_modules():
+        if not hasattr(module, "in_channels") or module.in_channels != 3:
+            continue
+
+        # Check if this is a wrapper with inner .conv attribute
+        if hasattr(module, "conv") and isinstance(module.conv, nn.Conv2d):
+            old_conv = module.conv
+            module.conv = adapt_input_channels(
+                old_conv, new_in_channels=1, pretrained=pretrained
+            )
+            adapted_count += 1
+
+        elif isinstance(module, nn.Conv2d):
+            # Direct Conv2d - need to replace it in parent
+            parts = name.split(".")
+            parent = backbone
+            for part in parts[:-1]:
+                if part.isdigit():
+                    parent = parent[int(part)]
+                else:
+                    parent = getattr(parent, part)
+
+            child_name = parts[-1]
+            new_conv = adapt_input_channels(
+                module, new_in_channels=1, pretrained=pretrained
+            )
+
+            if child_name.isdigit():
+                parent[int(child_name)] = new_conv
+            else:
+                setattr(parent, child_name, new_conv)
+
+            adapted_count += 1
+
+        if not adapt_all and adapted_count > 0:
+            break
+
+    return adapted_count

wavedl/models/base.py

@@ -15,6 +15,54 @@ import torch
 import torch.nn as nn
 
 
+# =============================================================================
+# TYPE ALIASES
+# =============================================================================
+
+# Spatial shape type aliases for model input dimensions
+SpatialShape1D = tuple[int]
+SpatialShape2D = tuple[int, int]
+SpatialShape3D = tuple[int, int, int]
+SpatialShape = SpatialShape1D | SpatialShape2D | SpatialShape3D
+
+
+# =============================================================================
+# UTILITY FUNCTIONS
+# =============================================================================
+
+
+def compute_num_groups(num_channels: int, preferred_groups: int = 32) -> int:
+    """
+    Compute valid num_groups for GroupNorm that divides num_channels evenly.
+
+    GroupNorm requires num_channels to be divisible by num_groups. This function
+    finds the largest valid divisor up to preferred_groups.
+
+    Args:
+        num_channels: Number of channels to normalize (must be positive)
+        preferred_groups: Preferred number of groups (default: 32)
+
+    Returns:
+        Valid num_groups that satisfies num_channels % num_groups == 0
+
+    Example:
+        >>> compute_num_groups(64)  # Returns 32
+        >>> compute_num_groups(48)  # Returns 16 (48 % 32 != 0)
+        >>> compute_num_groups(7)  # Returns 1 (prime number)
+    """
+    # Try preferred groups first, then common divisors
+    for groups in [preferred_groups, 16, 8, 4, 2, 1]:
+        if groups <= num_channels and num_channels % groups == 0:
+            return groups
+
+    # Fallback: find any valid divisor (always returns at least 1)
+    for groups in range(min(32, num_channels), 0, -1):
+        if num_channels % groups == 0:
+            return groups
+
+    return 1  # Always valid
+
+
 class BaseModel(nn.Module, ABC):
     """
     Abstract base class for all regression models.

wavedl/models/caformer.py

@@ -33,7 +33,7 @@ Author: Ductho Le (ductho.le@outlook.com)
 import torch
 import torch.nn as nn
 
-from wavedl.models._timm_utils import build_regression_head
+from wavedl.models._pretrained_utils import build_regression_head
 from wavedl.models.base import BaseModel
 from wavedl.models.registry import register_model
 

wavedl/models/cnn.py

@@ -24,14 +24,10 @@ from typing import Any
 import torch
 import torch.nn as nn
 
-from wavedl.models.base import BaseModel
+from wavedl.models.base import BaseModel, SpatialShape, compute_num_groups
 from wavedl.models.registry import register_model
 
 
-# Type alias for spatial shapes
-SpatialShape = tuple[int] | tuple[int, int] | tuple[int, int, int]
-
-
 def _get_conv_layers(
     dim: int,
 ) -> tuple[type[nn.Module], type[nn.Module], type[nn.Module]]:
@@ -163,27 +159,6 @@ class CNN(BaseModel):
             nn.Linear(64, out_size),
         )
 
-    @staticmethod
-    def _compute_num_groups(num_channels: int, target_groups: int = 4) -> int:
-        """
-        Compute valid num_groups for GroupNorm that divides num_channels.
-
-        Finds the largest divisor of num_channels that is <= target_groups,
-        or falls back to 1 if no suitable divisor exists.
-
-        Args:
-            num_channels: Number of channels (must be positive)
-            target_groups: Desired number of groups (default: 4)
-
-        Returns:
-            Valid num_groups that satisfies num_channels % num_groups == 0
-        """
-        # Try target_groups down to 1, return first valid divisor
-        for g in range(min(target_groups, num_channels), 0, -1):
-            if num_channels % g == 0:
-                return g
-        return 1  # Fallback (always valid)
-
     def _make_conv_block(
         self, in_channels: int, out_channels: int, dropout: float = 0.0
     ) -> nn.Sequential:
@@ -198,7 +173,7 @@ class CNN(BaseModel):
         Returns:
             Sequential block: Conv → GroupNorm → LeakyReLU → MaxPool [→ Dropout]
         """
-        num_groups = self._compute_num_groups(out_channels, target_groups=4)
+        num_groups = compute_num_groups(out_channels, preferred_groups=4)
 
         layers = [
             self._Conv(in_channels, out_channels, kernel_size=3, padding=1),

wavedl/models/convnext.py

@@ -28,14 +28,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from wavedl.models.base import BaseModel
+from wavedl.models.base import BaseModel, SpatialShape
 from wavedl.models.registry import register_model
 
 
-# Type alias for spatial shapes
-SpatialShape = tuple[int] | tuple[int, int] | tuple[int, int, int]
-
-
 def _get_conv_layer(dim: int) -> type[nn.Module]:
     """Get dimension-appropriate Conv class."""
     if dim == 1:
@@ -468,20 +464,11 @@ class ConvNeXtTinyPretrained(BaseModel):
         )
 
         # Modify first conv for single-channel input
-        old_conv = self.backbone.features[0][0]
-        self.backbone.features[0][0] = nn.Conv2d(
-            1,
-            old_conv.out_channels,
-            kernel_size=old_conv.kernel_size,
-            stride=old_conv.stride,
-            padding=old_conv.padding,
-            bias=old_conv.bias is not None,
+        from wavedl.models._pretrained_utils import adapt_first_conv_for_single_channel
+
+        adapt_first_conv_for_single_channel(
+            self.backbone, "features.0.0", pretrained=pretrained
         )
-        if pretrained:
-            with torch.no_grad():
-                self.backbone.features[0][0].weight = nn.Parameter(
-                    old_conv.weight.mean(dim=1, keepdim=True)
-                )
 
         if freeze_backbone:
             self._freeze_backbone()

wavedl/models/convnext_v2.py

@@ -31,20 +31,17 @@ from typing import Any
 import torch
 import torch.nn as nn
 
-from wavedl.models._timm_utils import (
+from wavedl.models._pretrained_utils import (
     LayerNormNd,
     build_regression_head,
     get_conv_layer,
     get_grn_layer,
     get_pool_layer,
 )
-from wavedl.models.base import BaseModel
+from wavedl.models.base import BaseModel, SpatialShape
 from wavedl.models.registry import register_model
 
 
-# Type alias for spatial shapes
-SpatialShape = tuple[int] | tuple[int, int] | tuple[int, int, int]
-
 __all__ = [
     "ConvNeXtV2Base",
     "ConvNeXtV2BaseLarge",
@@ -469,24 +466,11 @@ class ConvNeXtV2TinyPretrained(BaseModel):
 
     def _adapt_input_channels(self):
         """Adapt first conv layer for single-channel input."""
-        old_conv = self.backbone.features[0][0]
-        new_conv = nn.Conv2d(
-            1,
-            old_conv.out_channels,
-            kernel_size=old_conv.kernel_size,
-            stride=old_conv.stride,
-            padding=old_conv.padding,
-            bias=old_conv.bias is not None,
-        )
+        from wavedl.models._pretrained_utils import adapt_first_conv_for_single_channel
 
-        if self.pretrained:
-            with torch.no_grad():
-                # Average RGB weights for grayscale
-                new_conv.weight.copy_(old_conv.weight.mean(dim=1, keepdim=True))
-                if old_conv.bias is not None:
-                    new_conv.bias.copy_(old_conv.bias)
-
-        self.backbone.features[0][0] = new_conv
+        adapt_first_conv_for_single_channel(
+            self.backbone, "features.0.0", pretrained=self.pretrained
+        )
 
     def _freeze_backbone(self):
         """Freeze all backbone parameters except classifier."""

wavedl/models/densenet.py

@@ -26,14 +26,10 @@ from typing import Any
 import torch
 import torch.nn as nn
 
-from wavedl.models.base import BaseModel
+from wavedl.models.base import BaseModel, SpatialShape
 from wavedl.models.registry import register_model
 
 
-# Type alias for spatial shapes
-SpatialShape = tuple[int] | tuple[int, int] | tuple[int, int, int]
-
-
 def _get_layers(dim: int):
     """Get dimension-appropriate layer classes."""
     if dim == 1:
@@ -374,20 +370,11 @@ class DenseNet121Pretrained(BaseModel):
         )
 
         # Modify first conv for single-channel input
-        old_conv = self.backbone.features.conv0
-        self.backbone.features.conv0 = nn.Conv2d(
-            1,
-            old_conv.out_channels,
-            kernel_size=old_conv.kernel_size,
-            stride=old_conv.stride,
-            padding=old_conv.padding,
-            bias=False,
+        from wavedl.models._pretrained_utils import adapt_first_conv_for_single_channel
+
+        adapt_first_conv_for_single_channel(
+            self.backbone, "features.conv0", pretrained=pretrained
         )
-        if pretrained:
-            with torch.no_grad():
-                self.backbone.features.conv0.weight = nn.Parameter(
-                    old_conv.weight.mean(dim=1, keepdim=True)
-                )
 
         if freeze_backbone:
             self._freeze_backbone()