""" 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)