Spaces:

Shuxian
/

SGDNet-Fish-Classification

Running

SGDNet-Fish-Classification / models /sgd_net.py
a stable version
255fc2e 9 months ago
32.9 kB
"""
Shape-Guided Dual-Attention Network (SGD-Net)
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
import os
from pathlib import Path
from typing import List, Tuple, Optional


class AttentionGate(nn.Module):
        """Attention gate for skip connections in decoder"""
        
        def __init__(self, F_g, F_l, F_int):
            super(AttentionGate, self).__init__()
            self.W_g = nn.Sequential(
                nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),
                nn.BatchNorm2d(F_int)
            )
            
            self.W_x = nn.Sequential(
                nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),
                nn.BatchNorm2d(F_int)
            )
            
            self.psi = nn.Sequential(
                nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),
                nn.BatchNorm2d(1),
                nn.Sigmoid()
            )
            
            self.relu = nn.ReLU(inplace=True)
        
        def forward(self, g, x):
            g1 = self.W_g(g)
            x1 = self.W_x(x)
            
            # Handle size mismatch (upsample g1 to match x1 if needed)
            if g1.shape[2:] != x1.shape[2:]:
                g1 = F.interpolate(g1, size=x1.shape[2:], mode='bilinear', align_corners=False)
            
            psi = self.relu(g1 + x1)
            psi = self.psi(psi)
            return x * psi


class ShapeGuidedAttention(nn.Module):
        """
        Shape-Guided Attention module
        Uses decoder features to generate attention mask for RGB features
        """
        
        def __init__(self, rgb_channels, decoder_channels):
            super(ShapeGuidedAttention, self).__init__()
            self.decoder_conv = nn.Sequential(
                nn.Conv2d(decoder_channels, rgb_channels, kernel_size=1),
                nn.BatchNorm2d(rgb_channels),
                nn.Sigmoid()
            )
        
        def forward(self, rgb_feat, decoder_feat):
            """
            Args:
                rgb_feat: RGB features from backbone
                decoder_feat: Intermediate decoder features
            
            Returns:
                Enhanced RGB features with shape-guided attention
            """
            attention_mask = self.decoder_conv(decoder_feat)
            
            # Handle size mismatch (resize attention_mask to match rgb_feat if needed)
            if attention_mask.shape[2:] != rgb_feat.shape[2:]:
                attention_mask = F.interpolate(
                    attention_mask, 
                    size=rgb_feat.shape[2:], 
                    mode='bilinear', 
                    align_corners=False
                )
            
            enhanced_rgb = rgb_feat * attention_mask
            return enhanced_rgb


class DepthDecoder(nn.Module):
        """UNet-style decoder for depth prediction"""
        
        def __init__(self, encoder_channels, decoder_channels=[256, 128, 64, 32], use_attention_gate=True):
            super(DepthDecoder, self).__init__()
            self.decoder_channels = decoder_channels
            self.use_attention_gate = use_attention_gate
            
            # Reverse encoder channels for skip connections
            # Exclude first (input) layer, keep only skip connection features
            skip_channels = encoder_channels[::-1][1:]  # Exclude first (input) layer
            
            # Adjust decoder_channels to match available skip connections
            num_skip = len(skip_channels)
            num_decoder = len(decoder_channels)
            
            if num_decoder > num_skip:
                # If we have more decoder stages than skip connections, reduce decoder stages
                decoder_channels = decoder_channels[:num_skip]
                self.decoder_channels = decoder_channels
            
            # Decoder blocks
            self.decoder_blocks = nn.ModuleList()
            in_channels = encoder_channels[-1]
            
            for i, out_channels in enumerate(decoder_channels):
                # Upsample + conv (always first)
                self.decoder_blocks.append(
                    nn.Sequential(
                        nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),
                        nn.BatchNorm2d(out_channels),
                        nn.ReLU(inplace=True),
                        nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
                        nn.BatchNorm2d(out_channels),
                        nn.ReLU(inplace=True)
                    )
                )
                
                # Attention gate (if enabled and we have a corresponding skip connection)
                if self.use_attention_gate and i < len(skip_channels):
                    # Attention gate: F_g is upsampled feature (out_channels), F_l is skip (skip_channels[i])
                    self.decoder_blocks.append(
                        AttentionGate(F_g=out_channels, F_l=skip_channels[i], F_int=out_channels)
                    )
                    use_attention = True
                else:
                    use_attention = False
                
                in_channels = out_channels
            
            # Store whether to use attention for each stage
            self.use_attention = [self.use_attention_gate and i < len(skip_channels) for i in range(len(decoder_channels))]
            
            # Channel reduction layers for combining upsampled features with attended skip connections
            self.channel_reduce = nn.ModuleList()
            for i, out_ch in enumerate(decoder_channels):
                if i < len(skip_channels):
                    # Need to reduce: out_ch (upsampled) + skip_channels[i] (attended skip) -> out_ch
                    self.channel_reduce.append(
                        nn.Conv2d(out_ch + skip_channels[i], out_ch, kernel_size=1)
                    )
                else:
                    self.channel_reduce.append(None)
            
            # Final upsampling to input size (256x256)
            # Calculate how many more upsampling steps needed
            # Encoder outputs: [64x64, 32x32, 16x16, 8x8] for 256x256 input
            # After decoder: should be 64x64, need to upsample 4x to get 256x256
            final_upsample_factor = 4  # 64 -> 256 (4x)
            self.final_upsample = nn.Sequential(
                nn.ConvTranspose2d(decoder_channels[-1], decoder_channels[-1], kernel_size=4, stride=4),
                nn.BatchNorm2d(decoder_channels[-1]),
                nn.ReLU(inplace=True)
            )
            
            # Final depth prediction
            self.depth_head = nn.Sequential(
                nn.Conv2d(decoder_channels[-1], 1, kernel_size=1),
                nn.ReLU(inplace=True)  # Depth should be positive
            )
        
        def forward(self, encoder_features):
            """
            Args:
                encoder_features: List of features from encoder [f1, f2, f3, f4]
            
            Returns:
                depth: Predicted depth map (B, 1, H, W)
                intermediate_features: Intermediate decoder features for shape-guided attention
            """
            # Reverse for decoder (from deep to shallow)
            skip_features = encoder_features[::-1]
            x = skip_features[0]  # Deepest feature
            
            intermediate_features = []
            block_idx = 0
            
            # Decode with skip connections
            for i in range(len(self.decoder_channels)):
                # Upsample first
                x = self.decoder_blocks[block_idx](x)
                block_idx += 1
                
                # Apply attention gate to skip connection (if enabled and available)
                if self.use_attention[i] and (i + 1) < len(skip_features):
                    # Attention gate: x (upsampled) is g, skip_features[i+1] is x (skip)
                    skip_att = self.decoder_blocks[block_idx](x, skip_features[i + 1])
                    block_idx += 1
                    # Concatenate and reduce channels
                    x = torch.cat([x, skip_att], dim=1)
                    x = self.channel_reduce[i](x)
                elif not self.use_attention[i] and (i + 1) < len(skip_features):
                    # No attention gate: directly concatenate skip connection
                    x = torch.cat([x, skip_features[i + 1]], dim=1)
                    x = self.channel_reduce[i](x)
                
                intermediate_features.append(x)
            
            # Final upsampling to input size (256x256)
            x = self.final_upsample(x)
            
            # Ensure output is exactly 256x256 (in case of rounding errors)
            if x.shape[2] != 256 or x.shape[3] != 256:
                x = F.interpolate(x, size=(256, 256), mode='bilinear', align_corners=False)
            
            # Predict depth
            depth = self.depth_head(x)
            
            return depth, intermediate_features


class SGDNet(nn.Module):
        """
        Shape-Guided Dual-Attention Network
        """
        
        def __init__(self, backbone_name='convnextv2_tiny', num_classes=2, pretrained=True, pretrained_path=None, use_shape_attention=True, use_attention_gate=True):
            super(SGDNet, self).__init__()
            
            # Load backbone
            # If pretrained_path is provided, load from local file
            # Otherwise, use pretrained flag (downloads if True)
            if pretrained_path is not None:
                # Load model without pretrained weights first
                self.backbone = timm.create_model(
                    backbone_name,
                    pretrained=False,
                    features_only=True,
                    out_indices=(0, 1, 2, 3)
                )
                # Load pretrained weights from local file, or download if not found
                if not os.path.exists(pretrained_path):
                    print(f"Pretrained path not found: {pretrained_path}")
                    print("Attempting to download pretrained model...")
                    # Download and save as .pth format
                    try:
                        temp_model = timm.create_model(backbone_name, pretrained=True, features_only=True, out_indices=(0, 1, 2, 3))
                        # Save state dict as .pth file
                        pretrained_path_obj = Path(pretrained_path)
                        pretrained_path_obj.parent.mkdir(parents=True, exist_ok=True)
                        torch.save(temp_model.state_dict(), pretrained_path)
                        print(f"Downloaded and saved pretrained model to: {pretrained_path}")
                        del temp_model
                    except Exception as e:
                        print(f"Failed to download pretrained model: {e}")
                        raise FileNotFoundError(f"Pretrained model file not found and download failed: {pretrained_path}")
                
                if os.path.exists(pretrained_path):
                    # Check file extension first
                    file_ext = Path(pretrained_path).suffix.lower()
                    
                    if file_ext == '.safetensors':
                        # Load safetensors format (HuggingFace)
                        try:
                            from safetensors import safe_open
                            state_dict = {}
                            with safe_open(pretrained_path, framework="pt", device="cpu") as f:
                                for key in f.keys():
                                    state_dict[key] = f.get_tensor(key)
                            print(f"Loaded safetensors from: {pretrained_path}")
                        except ImportError:
                            raise ImportError("safetensors library is required to load .safetensors files. Install with: pip install safetensors")
                        except Exception as e:
                            raise RuntimeError(f"Failed to load safetensors file: {e}")
                    else:
                        # Load .pth or .pt format
                        try:
                            from timm.models import load_pretrained
                            load_pretrained(self.backbone, pretrained_path, strict=False)
                            print(f"Loaded pretrained weights via timm from: {pretrained_path}")
                            return  # Success, exit early
                        except:
                            # Fallback: manual loading with torch.load
                            checkpoint = torch.load(pretrained_path, map_location='cpu', weights_only=False)
                            # Handle different checkpoint formats
                            if isinstance(checkpoint, dict):
                                if 'state_dict' in checkpoint:
                                    state_dict = checkpoint['state_dict']
                                elif 'model' in checkpoint:
                                    state_dict = checkpoint['model']
                                else:
                                    state_dict = checkpoint
                            else:
                                state_dict = checkpoint
                    
                    # Filter and load only backbone weights (for safetensors or manual torch.load)
                    backbone_state_dict = {}
                    for k, v in state_dict.items():
                        # Remove 'backbone.' prefix if present
                        new_k = k.replace('backbone.', '')
                        # Skip non-backbone layers
                        if not any(x in new_k for x in ['depth_decoder', 'class_head', 'shape_attention']):
                            backbone_state_dict[new_k] = v
                    
                    missing_keys, unexpected_keys = self.backbone.load_state_dict(backbone_state_dict, strict=False)
                    if missing_keys:
                        print(f"Warning: Missing keys when loading pretrained: {len(missing_keys)} keys")
                    if unexpected_keys:
                        print(f"Warning: Unexpected keys when loading pretrained: {len(unexpected_keys)} keys")
                    print(f"Successfully loaded pretrained weights from: {pretrained_path}")
                else:
                    raise FileNotFoundError(f"Pretrained model file not found: {pretrained_path}")
            else:
                self.backbone = timm.create_model(
                    backbone_name,
                    pretrained=pretrained,
                    features_only=True,
                    out_indices=(0, 1, 2, 3)
                )
            
            # Get encoder channels
            with torch.no_grad():
                dummy_input = torch.randn(1, 3, 256, 256)
                features = self.backbone(dummy_input)
                encoder_channels = [f.shape[1] for f in features]
            
            # Depth decoder
            self.depth_decoder = DepthDecoder(encoder_channels, use_attention_gate=use_attention_gate)
            
            # Shape-guided attention (uses intermediate decoder features)
            self.use_shape_attention = use_shape_attention
            if use_shape_attention:
                # Attention for mid-level features (optional, for decoder guidance)
                self.shape_attention_mid = ShapeGuidedAttention(
                    rgb_channels=encoder_channels[2],  # Use mid-level RGB features
                    decoder_channels=self.depth_decoder.decoder_channels[-2]  # Second-to-last decoder channel
                )
                # Attention for deepest features (used for classification)
                self.shape_attention_deep = ShapeGuidedAttention(
                    rgb_channels=encoder_channels[-1],  # Use deepest RGB features
                    decoder_channels=self.depth_decoder.decoder_channels[-1]  # Last decoder channel
                )
            else:
                self.shape_attention_mid = None
                self.shape_attention_deep = None
            
            # Feature fusion: RGB + Depth
            # We'll use the last decoder intermediate feature (before final upsampling)
            # This has decoder_channels[-1] channels
            depth_feature_channels = self.depth_decoder.decoder_channels[-1]
            fused_channels = encoder_channels[-1] + depth_feature_channels
            
            # Classification head (now accepts fused features)
            self.class_head = nn.Sequential(
                nn.AdaptiveAvgPool2d(1),
                nn.Flatten(),
                nn.Linear(fused_channels, 256),
                nn.ReLU(inplace=True),
                nn.Dropout(0.5),
                nn.Linear(256, num_classes)
            )
        
        def forward(self, x):
            """
            Args:
                x: Input RGB image (B, 3, H, W)
            
            Returns:
                cls_logits: Classification logits (B, num_classes)
                depth: Predicted depth map (B, 1, H, W)
            """
            # Extract multi-scale features from backbone
            encoder_features = self.backbone(x)
            
            # Decode depth
            depth, decoder_intermediate = self.depth_decoder(encoder_features)
            
            # Get RGB features for classification (deepest level)
            rgb_features = encoder_features[-1]  # (B, C_rgb, H_rgb, W_rgb)
            
            # Get depth features for fusion (last intermediate decoder feature)
            depth_features = decoder_intermediate[-1]  # (B, C_depth, H_depth, W_depth)
            
            # Adapt depth features to match RGB feature spatial size
            # First, get spatial dimensions of RGB features
            B, C_rgb, H_rgb, W_rgb = rgb_features.shape
            _, C_depth, H_depth, W_depth = depth_features.shape
            
            # Resize depth features to match RGB feature spatial size
            if (H_depth, W_depth) != (H_rgb, W_rgb):
                depth_features = F.interpolate(
                    depth_features, 
                    size=(H_rgb, W_rgb), 
                    mode='bilinear', 
                    align_corners=False
                )
            
            # Apply shape-guided attention to RGB features (if enabled)
            if self.use_shape_attention:
                # Apply attention to mid-level features (helps decoder learn better depth)
                if self.shape_attention_mid is not None:
                    enhanced_rgb_mid = self.shape_attention_mid(
                        encoder_features[2],  # Mid-level RGB features
                        decoder_intermediate[-2]  # Second-to-last decoder feature
                    )
                
                # Apply attention to deepest features (used for classification)
                if self.shape_attention_deep is not None:
                    rgb_features = self.shape_attention_deep(
                        rgb_features,  # Deepest RGB features
                        depth_features  # Last decoder feature (already resized to match)
                    )
            
            # Feature fusion: concatenate RGB and depth features
            fused_features = torch.cat([rgb_features, depth_features], dim=1)  # (B, C_rgb + C_depth, H, W)
            
            # Classification using fused features
            cls_logits = self.class_head(fused_features)
            
            return cls_logits, depth


class ResNetSGDNet(nn.Module):
        """
        ResNet-based variant for ablation study
        """
        
        def __init__(self, backbone_name='resnet50', num_classes=2, pretrained=True, pretrained_path=None, use_shape_attention=True, use_attention_gate=True):
            super(ResNetSGDNet, self).__init__()
            
            # Load ResNet backbone
            # If pretrained_path is provided, load from local file
            if pretrained_path is not None:
                self.backbone = timm.create_model(
                    backbone_name,
                    pretrained=False,
                    features_only=True,
                    out_indices=(0, 1, 2, 3)
                )
                # Load pretrained weights from local file, or download if not found
                if not os.path.exists(pretrained_path):
                    print(f"Pretrained path not found: {pretrained_path}")
                    print("Attempting to download pretrained model...")
                    # Download and save as .pth format
                    try:
                        temp_model = timm.create_model(backbone_name, pretrained=True, features_only=True, out_indices=(0, 1, 2, 3))
                        # Save state dict as .pth file
                        pretrained_path_obj = Path(pretrained_path)
                        pretrained_path_obj.parent.mkdir(parents=True, exist_ok=True)
                        torch.save(temp_model.state_dict(), pretrained_path)
                        print(f"Downloaded and saved pretrained model to: {pretrained_path}")
                        del temp_model
                    except Exception as e:
                        print(f"Failed to download pretrained model: {e}")
                        raise FileNotFoundError(f"Pretrained model file not found and download failed: {pretrained_path}")
                
                if os.path.exists(pretrained_path):
                    # Check file extension
                    file_ext = Path(pretrained_path).suffix.lower()
                    
                    if file_ext == '.safetensors':
                        # Load safetensors format (HuggingFace)
                        try:
                            from safetensors import safe_open
                            state_dict = {}
                            with safe_open(pretrained_path, framework="pt", device="cpu") as f:
                                for key in f.keys():
                                    state_dict[key] = f.get_tensor(key)
                            print(f"Loaded safetensors from: {pretrained_path}")
                        except ImportError:
                            raise ImportError("safetensors library is required to load .safetensors files. Install with: pip install safetensors")
                        except Exception as e:
                            raise RuntimeError(f"Failed to load safetensors file: {e}")
                    else:
                        # Load .pth or .pt format
                        try:
                            from timm.models import load_pretrained
                            load_pretrained(self.backbone, pretrained_path, strict=False)
                            print(f"Loaded pretrained weights via timm from: {pretrained_path}")
                            return  # Success, exit early
                        except:
                            # Fallback: manual loading with torch.load
                            checkpoint = torch.load(pretrained_path, map_location='cpu', weights_only=False)
                            if isinstance(checkpoint, dict):
                                if 'state_dict' in checkpoint:
                                    state_dict = checkpoint['state_dict']
                                elif 'model' in checkpoint:
                                    state_dict = checkpoint['model']
                                else:
                                    state_dict = checkpoint
                            else:
                                state_dict = checkpoint
                    
                    # Filter and load only backbone weights
                    backbone_state_dict = {}
                    for k, v in state_dict.items():
                        new_k = k.replace('backbone.', '')
                        if not any(x in new_k for x in ['depth_decoder', 'class_head', 'shape_attention']):
                            backbone_state_dict[new_k] = v
                    
                    missing_keys, unexpected_keys = self.backbone.load_state_dict(backbone_state_dict, strict=False)
                    if missing_keys:
                        print(f"Warning: Missing keys when loading pretrained: {len(missing_keys)} keys")
                    if unexpected_keys:
                        print(f"Warning: Unexpected keys when loading pretrained: {len(unexpected_keys)} keys")
                    print(f"Successfully loaded pretrained weights from: {pretrained_path}")
                else:
                    raise FileNotFoundError(f"Pretrained model file not found: {pretrained_path}")
            else:
                self.backbone = timm.create_model(
                    backbone_name,
                    pretrained=pretrained,
                    features_only=True,
                    out_indices=(0, 1, 2, 3)
                )
            
            # Get encoder channels
            with torch.no_grad():
                dummy_input = torch.randn(1, 3, 256, 256)
                features = self.backbone(dummy_input)
                encoder_channels = [f.shape[1] for f in features]
            
            # Depth decoder
            self.depth_decoder = DepthDecoder(encoder_channels, use_attention_gate=use_attention_gate)
            
            # Shape-guided attention
            self.use_shape_attention = use_shape_attention
            if use_shape_attention:
                # Attention for mid-level features (optional, for decoder guidance)
                self.shape_attention_mid = ShapeGuidedAttention(
                    rgb_channels=encoder_channels[2],  # Use mid-level RGB features
                    decoder_channels=self.depth_decoder.decoder_channels[-2]  # Second-to-last decoder channel
                )
                # Attention for deepest features (used for classification)
                self.shape_attention_deep = ShapeGuidedAttention(
                    rgb_channels=encoder_channels[-1],  # Use deepest RGB features
                    decoder_channels=self.depth_decoder.decoder_channels[-1]  # Last decoder channel
                )
            else:
                self.shape_attention_mid = None
                self.shape_attention_deep = None
            
            # Feature fusion: RGB + Depth
            # We'll use the last decoder intermediate feature (before final upsampling)
            depth_feature_channels = self.depth_decoder.decoder_channels[-1]
            fused_channels = encoder_channels[-1] + depth_feature_channels
            
            # Classification head (now accepts fused features)
            self.class_head = nn.Sequential(
                nn.AdaptiveAvgPool2d(1),
                nn.Flatten(),
                nn.Linear(fused_channels, 256),
                nn.ReLU(inplace=True),
                nn.Dropout(0.5),
                nn.Linear(256, num_classes)
            )
        
        def forward(self, x):
            encoder_features = self.backbone(x)
            depth, decoder_intermediate = self.depth_decoder(encoder_features)
            
            # Get RGB features for classification (deepest level)
            rgb_features = encoder_features[-1]  # (B, C_rgb, H_rgb, W_rgb)
            
            # Get depth features for fusion (last intermediate decoder feature)
            depth_features = decoder_intermediate[-1]  # (B, C_depth, H_depth, W_depth)
            
            # Adapt depth features to match RGB feature spatial size
            # First, get spatial dimensions of RGB features
            B, C_rgb, H_rgb, W_rgb = rgb_features.shape
            _, C_depth, H_depth, W_depth = depth_features.shape
            
            # Resize depth features to match RGB feature spatial size
            if (H_depth, W_depth) != (H_rgb, W_rgb):
                depth_features = F.interpolate(
                    depth_features, 
                    size=(H_rgb, W_rgb), 
                    mode='bilinear', 
                    align_corners=False
                )
            
            # Apply shape-guided attention to RGB features (if enabled)
            if self.use_shape_attention:
                # Apply attention to mid-level features (helps decoder learn better depth)
                if self.shape_attention_mid is not None:
                    enhanced_rgb_mid = self.shape_attention_mid(
                        encoder_features[2],  # Mid-level RGB features
                        decoder_intermediate[-2]  # Second-to-last decoder feature
                    )
                
                # Apply attention to deepest features (used for classification)
                if self.shape_attention_deep is not None:
                    rgb_features = self.shape_attention_deep(
                        rgb_features,  # Deepest RGB features
                        depth_features  # Last decoder feature (already resized to match)
                    )
            
            # Feature fusion: concatenate RGB and depth features
            fused_features = torch.cat([rgb_features, depth_features], dim=1)  # (B, C_rgb + C_depth, H, W)
            
            # Classification using fused features
            cls_logits = self.class_head(fused_features)
            return cls_logits, depth


def get_timm_cache_dir():
        """
        Get the cache directory where timm stores pretrained models
        
        Returns:
            Path to timm cache directory
        """
        # Check TORCH_HOME environment variable first
        torch_home = os.environ.get('TORCH_HOME')
        if torch_home:
            cache_dir = Path(torch_home) / 'hub' / 'checkpoints'
        else:
            # Default location on Windows
            if os.name == 'nt':  # Windows
                cache_dir = Path.home() / '.cache' / 'torch' / 'hub' / 'checkpoints'
            else:  # Linux/Mac
                cache_dir = Path.home() / '.cache' / 'torch' / 'hub' / 'checkpoints'
        
        return cache_dir


def find_pretrained_model(backbone_name: str, cache_dir: Optional[Path] = None) -> Optional[str]:
        """
        Find pretrained model file in cache directory
        
        Args:
            backbone_name: Name of the backbone model (e.g., 'convnextv2_tiny', 'resnet50')
            cache_dir: Optional cache directory path. If None, uses default timm cache.
        
        Returns:
            Path to pretrained model file if found, None otherwise
        """
        if cache_dir is None:
            cache_dir = get_timm_cache_dir()
        
        # Common model file patterns
        patterns = [
            f"{backbone_name}*.pth",
            f"{backbone_name}*.pt",
            f"*{backbone_name}*.pth",
            f"*{backbone_name}*.pt",
        ]
        
        if not cache_dir.exists():
            return None
        
        # Search for model files
        for pattern in patterns:
            matches = list(cache_dir.glob(pattern))
            if matches:
                # Return the most recent file if multiple matches
                return str(max(matches, key=lambda p: p.stat().st_mtime))
        
        return None


def download_and_save_pretrained_model(backbone_name: str, save_dir: Optional[Path] = None) -> str:
        """
        Download pretrained model and save to specified directory
        
        Args:
            backbone_name: Name of the backbone model
            save_dir: Directory to save the model. If None, saves to timm cache.
        
        Returns:
            Path to saved model file
        """
        if save_dir is None:
            save_dir = get_timm_cache_dir()
        save_dir = Path(save_dir)
        save_dir.mkdir(parents=True, exist_ok=True)
        
        # Create model to trigger download
        model = timm.create_model(backbone_name, pretrained=True, features_only=False)
        
        # Get the default pretrained config to find the model URL/file
        from timm.models import safe_model_name
        safe_name = safe_model_name(backbone_name)
        
        # Try to find the downloaded file in cache
        cache_dir = get_timm_cache_dir()
        model_file = find_pretrained_model(backbone_name, cache_dir)
        
        if model_file:
            print(f"Model downloaded to: {model_file}")
            return model_file
        else:
            # If not found, save the state dict manually
            save_path = save_dir / f"{safe_name}_pretrained.pth"
            torch.save(model.state_dict(), save_path)
            print(f"Model saved to: {save_path}")
            return str(save_path)