a stable version

README.md

@@ -1,14 +1,63 @@
----
-title: SGDNet Fish Classification
-emoji: 🦀
-colorFrom: blue
-colorTo: purple
-sdk: gradio
-sdk_version: 6.5.1
-app_file: app.py
-pinned: false
-license: gpl-3.0
-short_description: SGD Network for robust fish classification
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SGD-Net: Shape-Guided Dual-Attention Network for Fish Classification
+
+Hugging Face Spaces demo for SGD-Net fish classification system.
+
+## Overview
+
+SGD-Net is a deep learning framework designed for robust fish classification (wild vs. farmed) under challenging environmental conditions. The method combines RGB and depth information through a novel shape-guided dual-attention mechanism, achieving superior performance and robustness compared to baseline methods.
+
+## Features
+
+- **Multi-modal Fusion**: Integrates RGB and depth information for enhanced feature representation
+- **Shape-Guided Attention**: Leverages depth-estimated shape information to guide RGB feature learning
+- **Dual-Attention Mechanism**: Combines shape-guided attention and attention gates for better feature fusion
+- **Real-time Inference**: Fast classification with pre-trained models
+- **Grad-CAM Visualization**: Visualize model attention with heatmaps
+
+## Usage
+
+Upload an image of a fish, and the system will:
+1. Detect if it's a yellow croaker (confidence threshold: 0.5)
+2. If detected, classify it as Wild or Farmed using SGD-Net
+3. Display detection mask and Grad-CAM heatmap for interpretability
+
+## Model Architecture
+
+- **Backbone**: ResNet50
+- **Depth Decoder**: UNet-style decoder with attention gates
+- **Classification Head**: Fully connected layers with dropout
+
+## Model Files
+
+- `sgd_net_checkpoint.pth`: Trained SGD-Net model weights
+
+## Installation
+
+```bash
+pip install -r requirements.txt
+```
+
+## Local Testing
+
+```bash
+python app.py
+```
+
+The application will start at `http://127.0.0.1:7860`
+
+## Citation
+
+If you use this code in your research, please cite:
+
+```bibtex
+@article{sgdnet2024,
+  title={SGD-Net: Shape-Guided Dual-Attention Network for Robust Fish Classification},
+  author={Your Name},
+  journal={Your Journal},
+  year={2024}
+}
+```
+
+## License
+
+See LICENSE file for details.

SETUP.md

@@ -0,0 +1,56 @@
+# Hugging Face Spaces 设置说明
+
+## 需要复制的模型文件
+
+在部署到 Hugging Face Spaces 之前，需要将以下模型文件复制到 `huggingface_space/` 目录：
+
+### 1. YOLO 模型
+- **源文件路径**: `E:\ECSF\dahuangyu\code\runs\segment\train4\weights\best.pt`
+- **目标文件名**: `huggingface_space/yolo_model.pt`
+- **说明**: YOLO分割模型，用于检测大黄鱼并生成mask
+
+### 2. SGD-Net 模型
+- **源文件路径**: `experiments\exp3_4_strategy_ours_stage2_finetune\checkpoint_final.pth`
+- **目标文件名**: `huggingface_space/sgd_net_checkpoint.pth`
+- **说明**: 训练好的SGD-Net分类模型权重
+
+## 复制命令（Windows PowerShell）
+
+```powershell
+# 复制YOLO模型
+Copy-Item "E:\ECSF\dahuangyu\code\runs\segment\train4\weights\best.pt" -Destination "huggingface_space\yolo_model.pt"
+
+# 复制SGD-Net模型
+Copy-Item "experiments\exp3_4_strategy_ours_stage2_finetune\checkpoint_final.pth" -Destination "huggingface_space\sgd_net_checkpoint.pth"
+```
+
+## 文件结构
+
+部署后的文件结构应该是：
+
+```
+huggingface_space/
+├── app.py                    # Gradio应用主文件
+├── gradcam.py                # Grad-CAM实现
+├── models/
+│   ├── __init__.py
+│   └── sgd_net.py           # SGD-Net模型定义
+├── requirements.txt          # 依赖
+├── README.md                # 说明文档
+├── SETUP.md                 # 设置说明（本文件）
+├── yolo_model.pt            # YOLO模型文件（需要复制）
+└── sgd_net_checkpoint.pth   # SGD-Net模型文件（需要复制）
+```
+
+## 部署到 Hugging Face Spaces
+
+1. 将整个 `huggingface_space/` 目录上传到 Hugging Face Spaces
+2. 确保所有模型文件都已上传
+3. Hugging Face Spaces 会自动安装 `requirements.txt` 中的依赖
+4. 应用会自动运行 `app.py`
+
+## 注意事项
+
+- 模型文件可能较大（几百MB），上传需要一些时间
+- 确保模型文件路径在代码中正确
+- 首次运行时，timm会自动下载ResNet50的预训练权重

app.py

@@ -0,0 +1,408 @@
+"""
+Hugging Face Spaces App for SGD-Net Fish Classification
+Enhanced with Grad-CAM visualization
+"""
+import gradio as gr
+import torch
+import torch.nn.functional as F
+import numpy as np
+import cv2
+from pathlib import Path
+from ultralytics import YOLO
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import matplotlib.pyplot as plt
+from matplotlib.colors import LinearSegmentedColormap
+
+from models.sgd_net import ResNetSGDNet
+from gradcam import GradCAM, get_target_layer
+
+# Model paths
+YOLO_MODEL_PATH = "yolo_model.pt"
+SGD_NET_CHECKPOINT = "sgd_net_checkpoint.pth"
+
+# SGD-Net configuration
+SGD_NET_CONFIG = {
+    "backbone": "resnet50",
+    "num_classes": 2,
+    "pretrained_path": None,
+    "use_shape_attention": False,
+    "use_attention_gate": True
+}
+
+# Global variables for models
+yolo_model = None
+sgd_net_model = None
+gradcam = None
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Preprocessing transform
+transform = A.Compose([
+    A.Resize(256, 256),
+    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ToTensorV2()
+])
+
+
+def load_models():
+    """Load YOLO and SGD-Net models"""
+    global yolo_model, sgd_net_model, gradcam
+    
+    # Load YOLO model
+    print("Loading YOLO model...")
+    try:
+        yolo_model = YOLO(YOLO_MODEL_PATH)
+        print("YOLO model loaded successfully")
+    except Exception as e:
+        print(f"Error loading YOLO model: {e}")
+        raise
+    
+    # Load SGD-Net model
+    print("Loading SGD-Net model...")
+    sgd_net_model = ResNetSGDNet(
+        backbone_name=SGD_NET_CONFIG["backbone"],
+        num_classes=SGD_NET_CONFIG["num_classes"],
+        pretrained=True,
+        pretrained_path=SGD_NET_CONFIG.get("pretrained_path"),
+        use_shape_attention=SGD_NET_CONFIG["use_shape_attention"],
+        use_attention_gate=SGD_NET_CONFIG["use_attention_gate"]
+    )
+    
+    # Load checkpoint
+    print(f"Loading checkpoint from: {SGD_NET_CHECKPOINT}")
+    checkpoint = torch.load(SGD_NET_CHECKPOINT, map_location=device, weights_only=False)
+    
+    # Handle different checkpoint formats
+    if isinstance(checkpoint, dict):
+        if 'model_state_dict' in checkpoint:
+            state_dict = checkpoint['model_state_dict']
+        elif 'state_dict' in checkpoint:
+            state_dict = checkpoint['state_dict']
+        else:
+            state_dict = checkpoint
+    else:
+        state_dict = checkpoint
+    
+    # Filter mismatched keys
+    model_state_dict = sgd_net_model.state_dict()
+    filtered_state_dict = {}
+    for k, v in state_dict.items():
+        if k in model_state_dict and model_state_dict[k].shape == v.shape:
+            filtered_state_dict[k] = v
+        elif k.startswith('module.'):
+            k_new = k[7:]
+            if k_new in model_state_dict and model_state_dict[k_new].shape == v.shape:
+                filtered_state_dict[k_new] = v
+    
+    sgd_net_model.load_state_dict(filtered_state_dict, strict=False)
+    sgd_net_model = sgd_net_model.to(device)
+    sgd_net_model.eval()
+    print("SGD-Net model loaded successfully")
+    
+    # Initialize Grad-CAM
+    print("Initializing Grad-CAM...")
+    target_layer = get_target_layer(sgd_net_model)
+    if target_layer is not None:
+        gradcam = GradCAM(sgd_net_model, target_layer)
+        print("Grad-CAM initialized successfully")
+    else:
+        print("Warning: Could not find target layer for Grad-CAM")
+
+
+def get_yolo_mask_and_confidence(rgb_image, conf_thres=0.5):
+    """
+    Generate mask using YOLO and return confidence
+    
+    Returns:
+        mask: Binary mask (H, W) uint8, 255=foreground, 0=background
+        confidence: Maximum confidence, returns 0.0 if no detection
+    """
+    results = yolo_model.predict(rgb_image, conf=conf_thres, verbose=False)
+    result = results[0]
+    
+    confidence = 0.0
+    mask = np.zeros(rgb_image.shape[:2], dtype=np.uint8)
+    
+    if result.masks is not None and len(result.masks) > 0:
+        boxes = result.boxes
+        if len(boxes) > 0:
+            confs = boxes.conf.cpu().numpy()
+            best_idx = int(np.argmax(confs))
+            confidence = float(confs[best_idx])
+            
+            mask_data = result.masks[best_idx].data.cpu().numpy()
+            if mask_data.ndim == 3:
+                mask_data = mask_data[0]
+            
+            mask = (mask_data * 255).astype(np.uint8)
+            
+            if mask.shape != rgb_image.shape[:2]:
+                mask = cv2.resize(mask, (rgb_image.shape[1], rgb_image.shape[0]), 
+                                interpolation=cv2.INTER_NEAREST)
+    
+    return mask, confidence
+
+
+def crop_image_by_mask(rgb_image, mask, expand_ratio=0.1):
+    """
+    Crop image and mask based on mask bounding box
+    
+    Args:
+        rgb_image: RGB image (H, W, 3) numpy array
+        mask: Binary mask (H, W) numpy array, 255=foreground, 0=background
+        expand_ratio: Expansion ratio for bounding box
+    
+    Returns:
+        rgb_cropped: Cropped RGB image
+        mask_cropped: Cropped mask
+    """
+    mask_binary = mask > 0
+    if mask_binary.sum() > 0:
+        coords = np.column_stack(np.where(mask_binary))
+        y_min, x_min = coords.min(axis=0)
+        y_max, x_max = coords.max(axis=0)
+        
+        H_img, W_img = rgb_image.shape[:2]
+        margin_y = int((y_max - y_min) * expand_ratio)
+        margin_x = int((x_max - x_min) * expand_ratio)
+        
+        y_min = max(0, y_min - margin_y)
+        y_max = min(H_img, y_max + margin_y)
+        x_min = max(0, x_min - margin_x)
+        x_max = min(W_img, x_max + margin_x)
+        
+        # Crop RGB and mask based on bounding box
+        rgb_cropped = rgb_image[y_min:y_max, x_min:x_max].copy()
+        mask_cropped = mask[y_min:y_max, x_min:x_max].copy()
+    else:
+        # If mask is empty, use full image
+        rgb_cropped = rgb_image
+        mask_cropped = mask
+    
+    return rgb_cropped, mask_cropped
+
+
+def preprocess_image(rgb_image, mask):
+    """
+    Preprocess image: crop by mask, resize and normalize
+    
+    Args:
+        rgb_image: RGB image (H, W, 3) numpy array
+        mask: Binary mask (H, W) numpy array
+    
+    Returns:
+        tensor: Preprocessed image tensor (1, 3, 256, 256)
+        rgb_cropped: Cropped RGB image (for visualization)
+        mask_cropped: Cropped mask (for visualization)
+    """
+    # Step 1: Crop image and mask based on mask bounding box
+    rgb_cropped, mask_cropped = crop_image_by_mask(rgb_image, mask, expand_ratio=0.1)
+    
+    # Step 2: Resize and normalize
+    transformed = transform(image=rgb_cropped)
+    rgb_tensor = transformed['image']
+    rgb_tensor = rgb_tensor.unsqueeze(0)
+    
+    return rgb_tensor, rgb_cropped, mask_cropped
+
+
+def generate_gradcam_heatmap(cam, original_image, mask):
+    """
+    Generate Grad-CAM heatmap overlaid on original image (only within mask)
+    
+    Args:
+        cam: CAM heatmap (H, W) numpy array, values in [0, 1]
+        original_image: Original image (H, W, 3) numpy array
+        mask: Binary mask (H, W) numpy array, 255=foreground, 0=background
+    
+    Returns:
+        overlay: Overlaid image (H, W, 3) numpy array
+    """
+    # Resize CAM to original image size
+    h, w = original_image.shape[:2]
+    cam_resized = cv2.resize(cam, (w, h), interpolation=cv2.INTER_LINEAR)
+    
+    # Resize mask to match image size if needed
+    if mask.shape[:2] != (h, w):
+        mask_resized = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
+    else:
+        mask_resized = mask
+    
+    # Normalize mask to [0, 1]
+    mask_normalized = (mask_resized > 0).astype(np.float32)
+    
+    # Apply mask to CAM (only show attention within mask)
+    cam_masked = cam_resized * mask_normalized
+    
+    # Create heatmap colormap (jet colormap)
+    cam_colored = plt.cm.jet(cam_masked)[:, :, :3]  # (H, W, 3) RGB
+    cam_colored = (cam_colored * 255).astype(np.uint8)
+    
+    # Overlay: original * (1 - alpha) + heatmap * alpha, but only within mask
+    alpha = 0.4
+    mask_3d = np.stack([mask_normalized] * 3, axis=2)  # (H, W, 3)
+    
+    overlay = (original_image.astype(np.float32) * (1 - alpha * mask_3d) + 
+               cam_colored.astype(np.float32) * (alpha * mask_3d)).astype(np.uint8)
+    
+    return overlay
+
+
+def classify_fish(image):
+    """
+    Classify fish image
+    
+    Args:
+        image: PIL Image or numpy array
+    
+    Returns:
+        tuple: (result_text, rgb_image, mask_image, gradcam_image)
+    """
+    if image is None:
+        return "Please upload an image", None, None, None
+    
+    # Convert to numpy array (RGB format)
+    if hasattr(image, 'convert'):
+        rgb_image = np.array(image.convert('RGB'))
+    else:
+        rgb_image = np.array(image)
+        if len(rgb_image.shape) == 3 and rgb_image.shape[2] == 4:  # RGBA
+            rgb_image = rgb_image[:, :, :3]
+    
+    # Step 1: YOLO detection
+    mask, confidence = get_yolo_mask_and_confidence(rgb_image, conf_thres=0.5)
+    
+    # Step 2: Check confidence
+    if confidence < 0.5:
+        result_text = f"**Detection Confidence: {confidence:.3f} < 0.5**\n\n**Result: Not Larimichthys crocea**"
+        mask_display = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
+        return result_text, rgb_image, mask_display, None
+    
+    # Step 3: Crop and preprocess image for classification
+    rgb_tensor, rgb_cropped, mask_cropped = preprocess_image(rgb_image, mask)
+    rgb_tensor = rgb_tensor.to(device)
+    
+    # Inference
+    with torch.no_grad():
+        cls_logits, _ = sgd_net_model(rgb_tensor)
+        probs = F.softmax(cls_logits, dim=1)
+        pred_class = cls_logits.argmax(dim=1).item()
+        prob_wild = probs[0, 0].item()
+        prob_farmed = probs[0, 1].item()
+    
+    # Generate Grad-CAM heatmap (only within mask)
+    gradcam_image = None
+    if gradcam is not None:
+        try:
+            cam = gradcam.generate_cam(rgb_tensor, class_idx=pred_class)
+            # Use cropped image and mask for Grad-CAM visualization
+            gradcam_image = generate_gradcam_heatmap(cam, rgb_cropped, mask_cropped)
+        except Exception as e:
+            print(f"Error generating Grad-CAM: {e}")
+            gradcam_image = None
+    
+    # Generate result text
+    class_name = "Wild" if pred_class == 0 else "Farmed"
+    
+    result_text = f"""**Detection Confidence: {confidence:.3f} ≥ 0.5**
+
+**Classification Result: {class_name}**
+
+**Probability Distribution:**
+- Wild: **{prob_wild:.1%}**
+- Farmed: **{prob_farmed:.1%}**"""
+    
+    # Convert mask to 3-channel for display (use cropped mask)
+    mask_display = cv2.cvtColor(mask_cropped, cv2.COLOR_GRAY2RGB)
+    
+    # Use cropped image for display
+    return result_text, rgb_cropped, mask_display, gradcam_image
+
+
+# Create Gradio interface
+def create_interface():
+    # Custom CSS styles
+    custom_css = """
+    .gradio-container {
+        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+    }
+    .main-header {
+        text-align: center;
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        padding: 20px;
+        border-radius: 10px;
+        color: white;
+        margin-bottom: 20px;
+    }
+    .result-box {
+        background: #f8f9fa;
+        border-radius: 8px;
+        padding: 15px;
+        border-left: 4px solid #667eea;
+    }
+    """
+    
+    with gr.Blocks(title="Illumination-Invariant Fish Classification", css=custom_css, theme=gr.themes.Soft()) as demo:
+        # Title
+        gr.Markdown("""
+        <div class="main-header">
+        <h1>Illumination-Invariant Identification of Wild and Farmed Larimichthys crocea</h1>
+        <p style="font-size: 14px; margin-top: 10px;">via Cross-Modal 3D Shape Priors</p>
+        </div>
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_input = gr.Image(
+                    label="Upload Image",
+                    type="numpy",
+                    height=300
+                )
+                classify_btn = gr.Button(
+                    "Classify",
+                    variant="primary",
+                    size="lg"
+                )
+            
+            with gr.Column(scale=2):
+                result_output = gr.Textbox(
+                    label="Classification Result",
+                    lines=8,
+                    elem_classes=["result-box"]
+                )
+        
+        with gr.Row():
+            with gr.Column():
+                image_output = gr.Image(
+                    label="Original Image",
+                    height=300
+                )
+            
+            with gr.Column():
+                mask_output = gr.Image(
+                    label="Detection Mask",
+                    height=300
+                )
+            
+            with gr.Column():
+                gradcam_output = gr.Image(
+                    label="Grad-CAM Heatmap",
+                    height=300
+                )
+        
+        classify_btn.click(
+            fn=classify_fish,
+            inputs=image_input,
+            outputs=[result_output, image_output, mask_output, gradcam_output]
+        )
+    
+    return demo
+
+
+if __name__ == "__main__":
+    # Load models
+    load_models()
+    
+    # Create and launch application
+    demo = create_interface()
+    demo.launch(share=False)

gradcam.py

@@ -0,0 +1,177 @@
+"""
+Grad-CAM implementation for visualization
+"""
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+
+class GradCAM:
+    """Grad-CAM实现"""
+    
+    def __init__(self, model, target_layer):
+        self.model = model
+        self.target_layer = target_layer
+        self.gradients = None
+        self.activations = None
+        self.forward_handle = None
+        self.backward_handle = None
+        
+        # 强制开启梯度（确保梯度能回传）
+        for param in self.model.parameters():
+            param.requires_grad = True
+        
+        # 注册hook
+        self.hook_layers()
+    
+    def hook_layers(self):
+        """注册forward和backward hook"""
+        if self.forward_handle is not None:
+            self.forward_handle.remove()
+        if self.backward_handle is not None:
+            self.backward_handle.remove()
+        
+        def backward_hook(module, grad_in, grad_out):
+            """标准的backward hook签名：(module, grad_in, grad_out)"""
+            if grad_out[0] is not None:
+                self.gradients = grad_out[0]
+        
+        def forward_hook(module, input, output):
+            """forward hook保存激活值"""
+            self.activations = output
+        
+        self.forward_handle = self.target_layer.register_forward_hook(forward_hook)
+        self.backward_handle = self.target_layer.register_backward_hook(backward_hook)
+    
+    def generate_cam(self, input_tensor, class_idx=None):
+        """生成CAM热力图"""
+        self.gradients = None
+        self.activations = None
+        
+        # 设置为train模式以允许梯度回传
+        self.model.train()
+        # 禁用dropout和batchnorm的随机性
+        for m in self.model.modules():
+            if isinstance(m, (torch.nn.BatchNorm2d, torch.nn.Dropout)):
+                m.eval()
+        
+        input_tensor = input_tensor.requires_grad_(True)
+        
+        # Forward pass
+        output = self.model(input_tensor)
+        
+        # 如果是SGD-Net，取分类输出
+        if isinstance(output, tuple):
+            output = output[0]
+        
+        # 如果没有指定类别，使用预测类别
+        if class_idx is None:
+            class_idx = output.argmax(dim=1).item()
+        
+        # Backward pass
+        self.model.zero_grad()
+        output[0, class_idx].backward(retain_graph=True)
+        
+        # 如果gradients仍然是None或全0，使用fallback方法
+        if self.gradients is None or (hasattr(self.gradients, 'abs') and self.gradients.abs().max().item() < 1e-8):
+            if self.activations is not None:
+                if len(self.activations.shape) == 4:
+                    activations = self.activations[0]  # (C, H, W)
+                else:
+                    activations = self.activations
+                cam = torch.mean(activations.abs(), dim=0)  # (H, W)
+                cam_min = cam.min()
+                cam_max = cam.max()
+                if cam_max > cam_min + 1e-8:
+                    cam = (cam - cam_min) / (cam_max - cam_min)
+                else:
+                    cam = torch.ones_like(cam) * 0.5
+                return cam.detach().cpu().numpy()
+        
+        if self.gradients is None or self.activations is None:
+            return np.ones((8, 8)) * 0.5
+        
+        # 计算权重（全局平均池化梯度）
+        if len(self.gradients.shape) == 4:
+            gradients = self.gradients[0]  # (C, H, W)
+        else:
+            gradients = self.gradients
+        
+        if len(self.activations.shape) == 4:
+            activations = self.activations[0]  # (C, H, W)
+        else:
+            activations = self.activations
+        
+        # 计算权重（对每个通道的梯度进行全局平均池化）
+        weights = torch.mean(gradients, dim=(1, 2), keepdim=True)  # (C, 1, 1)
+        
+        # 加权求和
+        cam = torch.sum(weights * activations, dim=0)  # (H, W)
+        cam = F.relu(cam)  # ReLU（只保留正梯度）
+        
+        # 移除左上角异常高值（针对SGD-Net）
+        h, w = cam.shape
+        corner_h = h // 3
+        corner_w = w // 3
+        corner_region = cam[:corner_h, :corner_w]
+        rest_region = torch.cat([
+            cam[corner_h:, :].flatten(),
+            cam[:, corner_w:].flatten()
+        ])
+        
+        if rest_region.numel() > 0:
+            rest_median = torch.median(rest_region)
+            rest_std = torch.std(rest_region)
+            corner_max = torch.max(corner_region)
+            corner_mean = torch.mean(corner_region)
+            
+            if corner_max > rest_median + 1.5 * rest_std or corner_mean > rest_median + 0.5 * rest_std:
+                for i in range(corner_h):
+                    for j in range(corner_w):
+                        dist_to_edge = min(i, j)
+                        weight = max(0.0, 1.0 - dist_to_edge / max(corner_h, corner_w) * 2.0)
+                        weight = min(1.0, weight)
+                        cam[i, j] = cam[i, j] * (1 - weight) + rest_median * weight
+        
+        # 归一化
+        cam_min = cam.min()
+        cam_max = cam.max()
+        
+        if cam_max > cam_min + 1e-8:
+            cam = (cam - cam_min) / (cam_max - cam_min)
+        else:
+            cam = torch.ones_like(cam) * 0.5
+        
+        cam = torch.clamp(cam, 0, 1)
+        
+        return cam.detach().cpu().numpy()
+    
+    def __del__(self):
+        """清理hook"""
+        if self.forward_handle is not None:
+            self.forward_handle.remove()
+        if self.backward_handle is not None:
+            self.backward_handle.remove()
+
+
+def get_target_layer(model):
+    """获取目标层用于Grad-CAM（针对ResNetSGDNet）"""
+    backbone = model.backbone
+    
+    # ResNet - 取layer4的最后一个block的最后一个conv
+    if hasattr(backbone, 'layer4'):
+        layer4 = backbone.layer4
+        if hasattr(layer4, '__len__') and len(layer4) > 0:
+            last_block = layer4[-1]
+            if hasattr(last_block, 'conv3'):
+                return last_block.conv3
+            elif hasattr(last_block, 'conv2'):
+                return last_block.conv2
+            else:
+                return last_block
+        else:
+            return layer4
+    else:
+        # 默认取最后一个模块
+        modules = list(backbone.children())
+        return modules[-1] if modules else None

models/__init__.py

@@ -0,0 +1 @@
+# Models package

models/sgd_net.py

@@ -0,0 +1,694 @@
+"""
+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)
+

requirements.txt

@@ -0,0 +1,10 @@
+torch>=1.12.0
+torchvision>=0.13.0
+timm>=0.6.0
+albumentations>=1.1.0
+opencv-python>=4.5.0
+numpy>=1.21.0
+ultralytics>=8.0.0
+gradio>=3.0.0
+Pillow>=9.0.0
+matplotlib>=3.5.0

sgd_net_checkpoint.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2bcba89c16de381ac7e137959436ecbc562057cf506d5419cac458e3b3b05f44
+size 125507694

yolo_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:33172ace0b55b8a3f34aef9585a4e6167dac4a3f45eeae164b8ea74edf4e814c
+size 5997028