clean release

README.md

@@ -1,3 +1,303 @@
----
-license: mit
----
+# CASWiT: Context-Aware Stage Wise Transformer for Ultra-High Resolution Semantic Segmentation
+
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+
+Official implementation of **CASWiT**, a dual-branch architecture for ultra-high resolution semantic segmentation that leverages cross-attention fusion between high-resolution and low-resolution branches.
+
+## 📋 Table of Contents
+
+- [Overview](#overview)
+- [Architecture](#architecture)
+- [Installation](#installation)
+- [Dataset Preparation](#dataset-preparation)
+- [Usage](#usage)
+- [Configuration](#configuration)
+- [Results](#results)
+- [Citation](#citation)
+- [License](#license)
+
+## 🎯 Overview
+
+CASWiT addresses the challenge of semantic segmentation on ultra-high resolution images by introducing a dual-branch architecture:
+
+- **HR Branch**: Processes high-resolution crops (512×512) for fine-grained detail
+- **LR Branch**: Processes low-resolution context (downsampled by 2×) for context
+- **Cross-Attention Fusion**: Enables HR features to attend to LR context at each encoder stage
+
+This design allows the model to capture both local details and global context, leading to improved segmentation performance on large-scale datasets.  
+In particular, CASWiT achieves **65.35 mIoU** on the **FLAIR-HUB** RGB-only UHR benchmark and **49.1 mIoU** on **URUR**, outperforming prior RGB/UHR baselines while remaining memory-efficient.
+
+## 🏗️ Architecture
+
+![CASWiT architecture](model/architecture.png)
+
+Key components:
+- **Dual Swin Transformer Backbones**: Two UPerNet-Swin encoders process HR and LR streams
+- **Cross-Attention Fusion Blocks**: Multi-head cross-attention at each encoder stage
+- **Auxiliary LR Supervision**: Additional supervision on LR branch for better training
+
+## 📦 Installation
+
+### Requirements
+
+- Python 3.8+
+- PyTorch 2.0+
+- CUDA 11.8+ (for GPU training)
+
+### Setup
+
+1. Clone the repository:
+```bash
+git clone https://github.com/yourusername/CASWiT.git
+cd CASWiT```
+
+2. Install dependencies:
+```bash
+pip install -r requirements.txt
+```
+
+## 📊 Dataset Preparation
+
+### FLAIR-HUB
+
+FlairHub is a large-scale ultra-high resolution semantic segmentation dataset. To prepare the dataset:
+
+1. Download the FlairHub dataset
+2. Run the preparation script to merge tiles:
+```bash
+python dataset/prepareFlairHub.py
+```
+
+The script will merge GeoTIFF tiles into larger mosaics suitable for training.
+
+### URUR
+
+URUR dataset should be organized as:
+```
+URUR/
+├── train/
+│   ├── image/
+│   └── label/
+├── val/
+│   ├── image/
+│   └── label/
+└── test/
+    ├── image/
+    └── label/
+```
+
+### SWISSIMAGE
+
+For SWISSIMAGE dataset:
+1. Download images using the provided CSV file:
+```bash
+python dataset/download_swissimage.py list_all_swiss_image_sept2025.csv
+```
+
+## 🚀 Usage
+
+### Training
+
+Train CASWiT on FlairHub:
+```bash
+python train/train.py configs/FlairHub.yaml
+```
+
+For distributed training (multi-GPU):
+```bash
+torchrun --nproc_per_node=4 train/train.py configs/FlairHub.yaml
+```
+
+### Evaluation
+
+Evaluate a trained model:
+```bash
+python train/eval.py configs/FlairHub.yaml weights/checkpoint.pth test
+```
+
+### Inference
+
+Run inference on a single image:
+```bash
+python train/inference.py configs/FlairHub.yaml weights/checkpoint.pth image.tif output.png
+```
+
+### Using Main Entry Point
+
+Alternatively, use the unified main script:
+```bash
+# Training
+python main.py train --config configs/FlairHub.yaml
+
+# Evaluation
+python main.py eval --config configs/FlairHub.yaml --checkpoint weights/checkpoint.pth
+
+# Inference a single image
+python main.py inference --config configs/FlairHub.yaml --checkpoint weights/checkpoint.pth --image image.tif --output pred.png
+```
+
+## ⚙️ Configuration
+
+Configuration files are in YAML format. Example structure:
+
+```yaml
+paths:
+  data_path: "/path/to/dataset"
+  dataset_name: ""
+  train_img_subdir: "train/img"
+  train_msk_subdir: "train/msk"
+  val_img_subdir: "val/img"
+  val_msk_subdir: "val/msk"
+  test_img_subdir: "test/img"
+  test_msk_subdir: "test/msk"
+  save_dir: "weights"
+  pretrained_path: ""
+
+model:
+  model_name: "openmmlab/upernet-swin-base"  # or swin-tiny, swin-large
+  num_classes: 15
+  cross_attention_heads: 1
+  fusion_mlp_ratio: 4.0
+  fusion_drop_path: 0.1
+  lr_supervision_weight: 0.5
+
+training:
+  batch_size: 4
+  num_workers: 8
+  num_epochs: 20
+  learning_rate: 0.00006
+  amp: true
+  seed: 1337
+  eta_min: 0.000001
+
+wandb:
+  use_wandb: true
+  project: "Fusion_HRLR"
+  entity: "your_entity"
+  run_name: "caswit_experiment"
+```
+
+### Key Parameters
+
+- `model_name`: Swin variant (`upernet-swin-tiny`, `upernet-swin-base`, `upernet-swin-large`)
+- `cross_attention_heads`: Number of attention heads in cross-attention blocks
+- `lr_supervision_weight`: Weight for LR branch auxiliary supervision
+
+## 📈 Results
+
+### FLAIR-HUB (RGB-only UHR protocol)
+
+We first evaluate CASWiT on the FLAIR-HUB ultra-high-resolution aerial benchmark under the RGB-only UHR protocol.
+
+| Model                             | mIoU (%) ↑ | mF1 (%) ↑ | mBIoU (%) ↑ |
+|----------------------------------|-----------:|----------:|------------:|
+| *RGB Baselines (official FLAIR-HUB)* |||| 
+| Swin-T + UPerNet                 |  62.01     |  75.27    |      –      |
+| Swin-S + UPerNet                 |  61.87     |  75.11    |      –      |
+| Swin-B + UPerNet                 |  64.05     |  76.88    |      –      |
+| Swin-B + UPerNet (retrained)     |  64.02     |  76.64    |    32.57    |
+| Swin-L + UPerNet                 |  63.36     |  76.35    |      –      |
+| *Ours (RGB-only UHR protocol)*   |||| 
+| **CASWiT-Base**                  |  65.11     |  77.71    |    35.87    |
+| **CASWiT-Base-SSL**              |**65.35**   |**77.87**  |  **35.99**  |
+
+CASWiT-Base already improves over the retrained Swin-B + UPerNet baseline, and CASWiT-Base-SSL further pushes performance to **65.35 mIoU** and **77.87 mF1**.  
+On mean boundary IoU, CASWiT-Base-SSL reaches **35.99 mBIoU**, which is a **+3.42 mBIoU** gain over the retrained Swin-B baseline (32.57).
+
+---
+
+### URUR
+
+We also evaluate CASWiT on the URUR ultra-high-resolution benchmark, comparing to both generic and UHR-specific segmentation models.
+
+| Model                                   | mIoU (%) ↑ | Mem (MB) ↓ |
+|----------------------------------------|-----------:|-----------:|
+| *Generic Models*                       |||
+| PSPNet                                 |  32.0      |   5482     |
+| ResNet18 + DeepLabv3+                  |  33.1      |   5508     |
+| STDC                                   |  42.0      |   7617     |
+| *UHR Models*                           |||
+| GLNet                                  |  41.2      |   3063     |
+| FCLt                                   |  43.1      |   4508     |
+| ISDNet                                 |  45.8      |   4920     |
+| WSDNet                                 |  46.9      |   4510     |
+| Boosting Dual-branch                   |  48.2      |   3682     |
+| **CASWiT-Base**                        |**48.7**    |   3530     |
+| **CASWiT-Base-SSL**                    |**49.1**    |   3530     |
+
+On URUR, CASWiT-Base already matches and slightly surpasses prior UHR-specific methods, and CASWiT-Base-SSL achieves **49.1 mIoU**, i.e. **+2.2 mIoU** over WSDNet and **+0.9 mIoU** over Boosting Dual-branch (UHRS), while remaining competitive in memory usage.
+
+
+## 🔬 Self-Supervised Learning
+
+CASWiT also supports self-supervised pre-training using SimMIM-style SSL (Simple Masked Image Modeling):
+
+```python
+from model.CASWiT_ssl import CASWiT_SSL
+
+model_ssl = CASWiT_SSL(
+    model_name="openmmlab/upernet-swin-base",
+    mask_ratio_hr=0.75,
+    mask_ratio_lr=0.5
+)
+```
+
+## 🛠️ Project Structure
+
+```
+CASWiT/
+├── model/
+│   ├── CASWiT.py          # Main model architecture
+│   └── CASWiT_ssl.py      # SSL variant
+├── dataset/
+│   ├── definition_dataset.py
+│   ├── download_swissimage.py
+│   └── prepareFlairHub.py
+├── configs/
+│   ├── FlairHub.yaml
+│   ├── URUR.yaml
+│   └── SWISSIMAGE.yaml
+├── utils/
+│   ├── metrics.py
+│   ├── logging.py
+│   └── attention_viz.py
+├── train/
+│   ├── train.py
+│   ├── eval.py
+│   └── inference.py
+├── weights/               # Model checkpoints
+├── main.py
+├── requirements.txt
+└── README.md
+```
+
+## 📝 Citation
+
+If you use CASWiT in your research, please cite:
+
+```bibtex
+@article{caswit2025,
+  title={CASWiT: Context-Aware Swin Transformer for Ultra-High Resolution Semantic Segmentation},
+  author={Masked for instance},
+  journal={},
+  year={2026}
+}
+```
+
+## 📄 License
+
+This project is licensed under the MIT License - see the LICENSE file for details.
+
+## 🙏 Acknowledgments
+
+- [UPerNet](https://github.com/open-mmlab/mmsegmentation) for the base segmentation architecture
+- [Swin Transformer](https://github.com/microsoft/Swin-Transformer) for the backbone
+- [FlairHub](https://github.com/IGNF/FlairHub) for the dataset
+- [URUR](https://github.com/jankyee/URUR) for the dataset
+
+## 📧 Contact
+
+For questions and issues, please open an issue on GitHub.
+
+---
+

configs/config_FlairHub.yaml

@@ -0,0 +1,34 @@
+paths:
+  data_path: "/mnt/Data/FlairHUB/data_flairhub/output/FLAIR1024_optimal"
+  dataset_name: "FLAIRHUB"
+  train_img_subdir: "train/img"
+  train_msk_subdir: "train/msk"
+  val_img_subdir: "valid/img"
+  val_msk_subdir: "valid/msk"
+  test_img_subdir: "test/img"
+  test_msk_subdir: "test/msk"
+  save_dir: "weights"
+  pretrained_path: "weights/CASWiT-Base-SSL_FLAIRHUB_15classes.pth" #
+model:
+  model_name: "openmmlab/upernet-swin-base"
+  num_classes: 15
+  cross_attention_heads: 1
+  ignore_index: 255
+  # Cross-fusion options
+  fusion_mlp_ratio: 4.0
+  fusion_drop_path: 0.1
+  lr_supervision_weight: 0.5
+training:
+  batch_size: 4
+  num_workers: 8
+  num_epochs: 20          
+  learning_rate: 0.00006
+  amp: true
+  seed: 42
+  eta_min: 0.000001
+wandb:
+  use_wandb: true
+  project: "CASWiT-Base"
+  entity: "soloo"
+  run_name: "CASWiT-Base_FLAIRHUB_1epoch"
+print_device: true

configs/config_SWISSIMAGE.yaml

@@ -0,0 +1,34 @@
+paths:
+  data_path: "/path/to/swissimage"
+  dataset_name: "SWISSIMAGE"
+  train_img_subdir: "train/img"
+  train_msk_subdir: "train/msk"
+  val_img_subdir: "val/img"
+  val_msk_subdir: "val/msk"
+  test_img_subdir: "test/img"
+  test_msk_subdir: "test/msk"
+  save_dir: "weights"
+  pretrained_path: ""
+model:
+  model_name: "openmmlab/upernet-swin-base"
+  num_classes: 15
+  cross_attention_heads: 1
+  ignore_index: 255
+  fusion_mlp_ratio: 4.0
+  fusion_drop_path: 0.1
+  lr_supervision_weight: 0.5
+training:
+  batch_size: 4
+  num_workers: 8
+  num_epochs: 20
+  learning_rate: 0.00006
+  amp: true
+  seed: 1337
+  eta_min: 0.000001
+wandb:
+  use_wandb: true
+  project: "Fusion_HRLR"
+  entity: "your_entity"
+  run_name: "swissimage_swin_base_swin_base_fusion_end_lr_supervised_1heads"
+print_device: true
+

configs/config_URUR.yaml

@@ -0,0 +1,35 @@
+paths:
+  data_path: "/mnt/Data/URUR"
+  dataset_name: "URUR"
+  train_img_subdir: "train/image"
+  train_msk_subdir: "train/label"
+  val_img_subdir: "val/image"
+  val_msk_subdir: "val/label"
+  test_img_subdir: "test/image"
+  test_msk_subdir: "test/label"
+  save_dir: "weights"
+  #pretrained_path: "weights/URUR_Last_8classes_alltrain_fusionend_hrlr_swinbase_swinbase_lrsupervised_pos_bias_none_12_epoch_15_head1.pth" # weights/fusion_hrlr_swintiny_swintiny_pos_bias_none_4_epoch_14_head1.pth
+  pretrained_path: "weights/from_mae_URUR_Lastv2_8classes_alltrain_fusionall_hrlr_swinbase_swinbase_lrsupervised_pos_bias_none_5_epoch_15_head1.pth"
+model:
+  model_name: "openmmlab/upernet-swin-base"
+  num_classes: 8
+  cross_attention_heads: 1
+  ignore_index: 255
+  # Cross-fusion options
+  fusion_mlp_ratio: 4.0
+  fusion_drop_path: 0.1
+  lr_supervision_weight: 0.5
+training:
+  batch_size: 5
+  num_workers: 8
+  num_epochs: 20            
+  learning_rate: 0.00006
+  amp: true
+  seed: 1337
+  eta_min: 0.000001
+wandb:
+  use_wandb: true
+  project: "Fusion_HRLR"
+  entity: "soloo"
+  run_name: "URUR_swin_base_swin_base_fusion_end_lr_supervised_1heads"
+print_device: true

dataset/__init__.py

@@ -0,0 +1,16 @@
+"""
+Dataset loaders and utilities.
+"""
+
+from dataset.definition_dataset import (
+    SemanticSegmentationDatasetFusion,
+    SemanticSegmentationDatasetHR,
+    build_transforms
+)
+
+__all__ = [
+    'SemanticSegmentationDatasetFusion',
+    'SemanticSegmentationDatasetHR',
+    'build_transforms',
+]
+

dataset/definition_dataset.py

@@ -0,0 +1,402 @@
+"""
+Dataset definitions for CASWiT training and evaluation.
+
+This module provides dataset classes for semantic segmentation with
+HR/LR dual-branch processing.
+"""
+
+import os
+import math
+from pathlib import Path
+from typing import Optional, Union, Tuple, List, Dict
+import numpy as np
+import torch
+from torch import Tensor
+from torch.utils.data import Dataset
+from PIL import Image
+from tifffile import imread as tiff_imread
+from torchvision import transforms
+
+
+class SemanticSegmentationDatasetFusion(Dataset):
+    """
+    Dataset for HR/LR fusion training on FLAIRHub.
+    
+    Returns (image_hr, mask_hr, image_lr, mask_lr):
+    - image_hr: 512x512 crop starting at (256, 256)
+    - image_lr: full image downsampled by factor 2
+    - mask >=15 replaced by 255 (ignore)
+    - transforms applied to images (ToTensor + Normalize) and mask -> LongTensor
+    """
+    def __init__(self, image_dir: Path, mask_dir: Path, transform: Optional[transforms.Compose] = None):
+        self.image_dir = Path(image_dir)
+        self.mask_dir = Path(mask_dir)
+        self.image_filenames = sorted(os.listdir(self.image_dir))
+        self.mask_filenames = sorted(os.listdir(self.mask_dir))
+        assert len(self.image_filenames) == len(self.mask_filenames), "Images/Masks count mismatch"
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.image_filenames)
+
+    def __getitem__(self, idx):
+        image_path = self.image_dir / self.image_filenames[idx]
+        mask_path = self.mask_dir / self.mask_filenames[idx]
+
+        image = load_image(image_path)
+        mask = load_mask(mask_path)
+        mask[mask >= 15] = 255 
+
+        # Crop HR at 512x512
+        hr_crop_size = 512
+        crop_x, crop_y = 256, 256
+
+        image_hr = image[crop_x:crop_x + hr_crop_size, crop_y:crop_y + hr_crop_size]
+        mask_hr = mask[crop_x:crop_x + hr_crop_size, crop_y:crop_y + hr_crop_size]
+        
+        # Downsample LR
+        image_lr = image[::2,::2,:]
+        mask_lr = mask[::2,::2]
+
+        if self.transform:
+            image_hr = self.transform(to_pil_uint8(image_hr))
+            image_lr = self.transform(to_pil_uint8(image_lr))
+        else:
+            image_hr = to_tensor_img(image_hr)
+            image_lr = to_tensor_img(image_lr)
+        mask_hr = torch.tensor(mask_hr, dtype=torch.long)
+        mask_lr = torch.tensor(mask_lr, dtype=torch.long)
+
+        return image_hr, mask_hr, image_lr, mask_lr
+
+
+class SemanticSegmentationDatasetHR(Dataset):
+    """
+    Dataset for HR-only training (single branch, no LR).
+    
+    Returns (image_hr, mask_hr):
+    - image_hr: 512x512 crop starting at (256, 256)
+    - mask >=15 replaced by 255 (ignore)
+    """
+    def __init__(self, image_dir: Path, mask_dir: Path, transform: Optional[transforms.Compose] = None):
+        self.image_dir = Path(image_dir)
+        self.mask_dir = Path(mask_dir)
+        self.image_filenames = sorted(os.listdir(self.image_dir))
+        self.mask_filenames = sorted(os.listdir(self.mask_dir))
+        assert len(self.image_filenames) == len(self.mask_filenames), "Images/Masks count mismatch"
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.image_filenames)
+
+    def __getitem__(self, idx):
+        image_path = self.image_dir / self.image_filenames[idx]
+        mask_path = self.mask_dir / self.mask_filenames[idx]
+
+        image = load_image(image_path)
+        mask = load_mask(mask_path)
+        mask[mask >= 15] = 255
+
+        crop_x, crop_y = 256, 256
+        image_hr = image[crop_x:crop_x + 512, crop_y:crop_y + 512]
+        mask_hr = mask[crop_x:crop_x + 512, crop_y:crop_y + 512]
+
+        if self.transform:
+            image_hr = self.transform(to_pil_uint8(image_hr))
+        else:
+            image_hr = to_tensor_img(image_hr)
+        mask_hr = torch.tensor(mask_hr, dtype=torch.long)
+        return image_hr, mask_hr
+
+
+# ----------------------------
+# Image Loading Functions for slifing windows without overlap on URUR, deepglobe and INRIA
+# ----------------------------
+
+def load_image(path: Union[str, Path]) -> np.ndarray:
+    """
+    Load an image as HxWx3 (RGB), float32 [0,1].
+    Handles both TIFF and PNG files gracefully.
+    """
+    p = str(path)
+    arr = None
+
+    # 1) Try TIFF first if available
+    if tiff_imread is not None:
+        try:
+            arr = tiff_imread(p)
+        except Exception:
+            arr = None
+
+    # 2) Fallback to PIL
+    if arr is None:
+        with Image.open(p) as im:
+            arr = np.array(im.convert("RGB"))  # HWC uint8
+
+    # Ensure HWC format
+    if arr.ndim == 2:
+        arr = np.stack((arr, arr, arr), axis=-1)  # HWC
+    elif arr.ndim == 3 and arr.shape[0] in (3, 4) and arr.shape[-1] not in (3, 4):
+        arr = np.moveaxis(arr, 0, -1)  # CHW -> HWC
+
+    # Keep 3 channels
+    c = arr.shape[-1]
+    if c == 4:
+        arr = arr[..., :3]
+    elif c == 1:
+        arr = np.repeat(arr, 3, axis=-1)
+
+    # Normalize -> float32 [0,1]
+    if arr.dtype is np.dtype(np.uint8):
+        arr = arr.astype(np.float32) / 255.0
+    else:
+        arr = arr.astype(np.float32, copy=False)
+        m = arr.max()
+        if m > 1.0:
+            arr = arr / m
+
+    return arr  # float32 HWC in [0,1]
+
+
+def load_mask(path: Union[str, Path]) -> np.ndarray:
+    """Load a mask as HxW int64 (labels). Handles both TIFF and PNG files."""
+    p = str(path)
+    m = None
+    if tiff_imread is not None:
+        try:
+            m = tiff_imread(p)
+        except Exception:
+            m = None
+    if m is None:
+        with Image.open(p) as im:
+            m = np.array(im)
+
+    # Force 2D
+    if m.ndim == 3:
+        m = m[..., 0]
+    return m.astype(np.int64, copy=False)
+
+
+# ----------------------------
+# Helper Functions
+# ----------------------------
+
+def crop_with_pad(img: np.ndarray, y0: int, x0: int, h: int, w: int, pad_val=0) -> np.ndarray:
+    """Extract a crop HxW with padding if necessary (img HxW[ xC])."""
+    H, W = img.shape[:2]
+    y1, x1 = y0 + h, x0 + w
+
+    pad_top   = max(0, -y0);    ys = max(0, y0)
+    pad_left  = max(0, -x0);    xs = max(0, x0)
+    pad_bot   = max(0, y1 - H); ye = min(H, y1)
+    pad_right = max(0, x1 - W); xe = min(W, x1)
+
+    sl = img[ys:ye, xs:xe]
+    pad_cfg = ((pad_top, pad_bot), (pad_left, pad_right)) + (((0, 0),) if img.ndim == 3 else ())
+    return np.pad(sl, pad_cfg, mode="constant", constant_values=pad_val)
+
+
+def resize_np_img(img_hwc_float01: np.ndarray, size_hw: Tuple[int, int]) -> np.ndarray:
+    """Resize HWC float32[0,1] -> HWC float32[0,1] using bilinear interpolation."""
+    Ht, Wt = size_hw
+    im = Image.fromarray((np.clip(img_hwc_float01, 0.0, 1.0) * 255.0).astype(np.uint8))
+    im = im.resize((Wt, Ht), resample=Image.BILINEAR)
+    out = np.asarray(im, dtype=np.uint8).astype(np.float32) / 255.0
+    if out.ndim == 2:
+        out = np.stack((out, out, out), axis=-1)
+    return out
+
+
+def resize_np_mask(mask_hw_int: np.ndarray, size_hw: Tuple[int, int]) -> np.ndarray:
+    """Resize mask HW using nearest neighbor via PIL, output int64."""
+    Ht, Wt = size_hw
+    mask = np.ascontiguousarray(mask_hw_int)
+    if mask.ndim != 2:
+        raise ValueError(f"resize_np_mask expects 2D mask HW, received shape={mask.shape}")
+
+    dt = mask.dtype
+    if dt in (np.int64, np.int32, np.int16, np.int8, np.uint16):
+        pil_arr, pil_mode = mask.astype(np.int32, copy=False), "I"
+    elif dt == np.uint8:
+        pil_arr, pil_mode = mask, "L"
+    else:
+        pil_arr, pil_mode = mask.astype(np.int32, copy=False), "I"
+
+    im = Image.fromarray(pil_arr, mode=pil_mode).resize((Wt, Ht), resample=Image.NEAREST)
+    return np.asarray(im).astype(np.int64, copy=False)
+
+
+def to_tensor_img(x: np.ndarray) -> Tensor:
+    """Convert HWC float32[0,1] -> CHW float32[0,1]."""
+    return torch.from_numpy(np.transpose(x, (2, 0, 1)).copy())
+
+
+def to_pil_uint8(img_float01_hwc: np.ndarray) -> Image.Image:
+    """Convert HWC float32[0,1] -> PIL RGB uint8."""
+    arr = (np.clip(img_float01_hwc, 0.0, 1.0) * 255.0).round().astype(np.uint8)
+    return Image.fromarray(arr, mode="RGB")
+
+
+# ----------------------------
+# Dataset Classes
+# ----------------------------
+
+class URURHRLRDataset(Dataset):
+    """
+    URUR dataset with HR/LR dual-branch processing and tiling support.
+    
+    In test mode, each worker caches the current image+mask in RAM for all tiles
+    of the same image to avoid re-reading for each tile.
+    """
+    def __init__(
+        self,
+        image_dir: Union[str, Path],
+        mask_dir: Union[str, Path],
+        num_classes: int,
+        mode: str = "train",
+        ignore_index: int = 255,
+        hr_size: int = 1024,
+        lr_side: int = 2048,
+        transform: Optional = None,
+        limit: Optional[int] = None
+    ) -> None:
+        assert mode in {"train", "val", "test"}
+        self.image_dir = Path(image_dir)
+        self.mask_dir = Path(mask_dir)
+        self.mode = mode
+        self.num_classes = int(num_classes)
+        self.ignore_index = int(ignore_index)
+        self.HR = int(hr_size)
+        self.LR_WIN = int(lr_side)
+        self.transform = transform
+
+        imgs = sorted([p for p in self.image_dir.iterdir() if p.is_file()])
+        msks = sorted([p for p in self.mask_dir.iterdir() if p.is_file()])
+        if limit is not None:
+            imgs, msks = imgs[:limit], msks[:limit]
+        assert len(imgs) == len(msks) and len(imgs) > 0, "Images/Masks missing or misaligned"
+
+        self.images: List[Path] = imgs
+        self.masks:  List[Path] = msks
+
+        # Tile index + quick sizes (without loading full image)
+        self._test_index: List[Tuple[int, int, int]] = []   # (img_id, y0, x0)
+        self._sizes: List[Tuple[int, int]] = []             # (H, W) per image
+
+        if self.mode == "test" or self.mode == "val":
+            for img_id, ip in enumerate(self.images):
+                with Image.open(ip) as im:
+                    W, H = im.size
+                self._sizes.append((H, W))
+
+                n_ty = math.ceil(H / self.HR)
+                n_tx = math.ceil(W / self.HR)
+                for iy in range(n_ty):
+                    for ix in range(n_tx):
+                        self._test_index.append((img_id, iy * self.HR, ix * self.HR))
+
+        # Cache per worker (used only in test mode)
+        self._cache_img_id: Optional[int] = None
+        self._cache_img: Optional[np.ndarray] = None   # float32 HWC [0,1]
+        self._cache_msk: Optional[np.ndarray] = None   # int64 HW
+
+    def __len__(self) -> int:
+        return len(self._test_index) if self.mode == "test" else len(self.images)
+
+    def _extract_pair_np(
+        self, img: np.ndarray, msk: np.ndarray, y0: int, x0: int
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        """Extract HR/LR pair in numpy (images HWC float32[0,1], masks HW int64)."""
+        # HR
+        img_hr = crop_with_pad(img, y0, x0, self.HR, self.HR, pad_val=0)
+        msk_hr = crop_with_pad(msk, y0, x0, self.HR, self.HR, pad_val=self.ignore_index)
+        # Resize HR to half size
+        img_hr = resize_np_img(img_hr, (self.HR//2, self.HR//2))
+
+        # LR centered on HR
+        cy, cx = y0 + self.HR // 2, x0 + self.HR // 2
+        half = self.LR_WIN // 2
+        img_lr = crop_with_pad(img, cy - half, cx - half, self.LR_WIN, self.LR_WIN, pad_val=0)
+        msk_lr = crop_with_pad(msk, cy - half, cx - half, self.LR_WIN, self.LR_WIN, pad_val=self.ignore_index)
+
+        # Downsample LR -> HR
+        img_lr_512 = resize_np_img(img_lr, (self.HR//2, self.HR//2))
+        msk_lr_512 = resize_np_mask(msk_lr, (self.HR, self.HR))
+
+        # Clamp out of range -> ignore_index
+        msk_hr = msk_hr.astype(np.int64, copy=False)
+        msk_lr_512 = msk_lr_512.astype(np.int64, copy=False)
+        msk_hr[msk_hr >= self.num_classes] = self.ignore_index
+        msk_lr_512[msk_lr_512 >= self.num_classes] = self.ignore_index
+
+        return img_hr, msk_hr, img_lr_512, msk_lr_512
+
+    def __getitem__(self, idx: int):
+        if self.mode == "test" or self.mode == "val":
+            # Tile index
+            img_id, y0, x0 = self._test_index[idx]
+            ip, mp = self.images[img_id], self.masks[img_id]
+            H, W = self._sizes[img_id]
+
+            # Cache per worker: read/convert only once per image
+            if self._cache_img_id != img_id:
+                self._cache_img = load_image(ip)   # float32 HWC [0,1]
+                self._cache_msk = load_mask(mp)    # int64 HW
+                self._cache_img_id = img_id
+
+            img = self._cache_img
+            msk = self._cache_msk
+
+            img_hr_np, msk_hr_np, img_lr_np, msk_lr_np = self._extract_pair_np(img, msk, y0, x0)
+
+            if self.transform:
+                image_hr = self.transform(to_pil_uint8(img_hr_np))
+                image_lr = self.transform(to_pil_uint8(img_lr_np))
+            else:
+                image_hr = to_tensor_img(img_hr_np)
+                image_lr = to_tensor_img(img_lr_np)
+
+            mask_hr = torch.as_tensor(msk_hr_np, dtype=torch.long)
+            mask_lr = torch.as_tensor(msk_lr_np, dtype=torch.long)
+
+            meta: Dict[str, object] = {
+                "img_path": str(ip),
+                "mask_path": str(mp),
+                "tile": (int(y0), int(x0), self.HR, self.HR),
+                "img_hw": (int(H), int(W)),
+                "tile_index": int(idx),
+            }
+            return image_hr, mask_hr, image_lr, mask_lr, meta
+
+        # train mode
+        ip, mp = self.images[idx], self.masks[idx]
+        img = load_image(ip)
+        msk = load_mask(mp)
+        H, W = img.shape[:2]
+
+        y0 = 0 if H <= self.HR else np.random.randint(0, H - self.HR + 1)
+        x0 = 0 if W <= self.HR else np.random.randint(0, W - self.HR + 1)
+
+        img_hr_np, msk_hr_np, img_lr_np, msk_lr_np = self._extract_pair_np(img, msk, y0, x0)
+
+        if self.transform:
+            image_hr = self.transform(to_pil_uint8(img_hr_np))
+            image_lr = self.transform(to_pil_uint8(img_lr_np))
+        else:
+            image_hr = to_tensor_img(img_hr_np)
+            image_lr = to_tensor_img(img_lr_np)
+
+        mask_hr = torch.as_tensor(msk_hr_np, dtype=torch.long)
+        mask_lr = torch.as_tensor(msk_lr_np, dtype=torch.long)
+
+        meta: Dict[str, object] = {"img_path": str(ip), "mask_path": str(mp)}
+        return image_hr, mask_hr, image_lr, mask_lr, meta
+
+
+def build_transforms():
+    """Build standard transforms with normalization (mean=std=0.5)."""
+    return transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    ])
+

dataset/download_swissimage.py

@@ -0,0 +1,67 @@
+"""
+SwissImage dataset downloader.
+
+Downloads SwissImage dataset from URLs provided in a CSV file.
+"""
+
+import os
+import sys
+import pandas as pd
+import requests
+
+
+def download_data(url: str, destination: str, count: int) -> int:
+    """
+    Download a single image from URL.
+    
+    Args:
+        url: Image URL to download
+        destination: Local file path to save
+        count: Current download count
+        
+    Returns:
+        Updated download count
+    """
+    response = requests.get(url)
+    if response.status_code == 200:
+        with open(destination, 'wb') as f:
+            f.write(response.content)
+        count += 1
+    else:
+        print(f"Failed to download: {url}. Status: {response.status_code}")
+        with open('failed_images.txt', 'a') as f:
+            f.write(url + '\n')
+    return count
+
+
+def main():
+    """Main function to download SwissImage dataset."""
+    if len(sys.argv) < 2:
+        print("Usage: python download_swissimage.py <csv_file>")
+        sys.exit(1)
+    
+    csv_file = sys.argv[1]
+    df = pd.read_csv(csv_file, header=None)
+    count_download = 0
+    count = 0
+    total = len(df)
+    print(f'Downloading {total} images.')
+    
+    if not os.path.exists('data'):
+        os.mkdir('data')
+    
+    for row in df.itertuples():
+        download_link = row[1]
+        count += 1
+        if count % 10 == 0:
+            print(f'Progress: {count/total*100:.1f}%')
+        fn = download_link.split('/')[-1]
+        fn_local = os.path.join('data', fn)
+        count_download = download_data(download_link, fn_local, count_download)
+    
+    print(f'Process finished with {count_download} images downloaded out of {total} planned')
+
+
+if __name__ == "__main__":
+    main()
+

dataset/prepareFlairHub.py

@@ -0,0 +1,270 @@
+"""
+FlairHub dataset preparation script.
+
+Merges GeoTIFF tiles from FlairHub into larger mosaics for training.
+This script processes the hierarchical folder structure and merges
+neighboring tiles into 2x2 mosaics.
+"""
+
+import rasterio
+from rasterio.windows import Window
+import os
+from affine import Affine
+import numpy as np
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor
+
+
+def get_raster_info(filepath):
+    """
+    Get raster information from a GeoTIFF file.
+    
+    Args:
+        filepath: Path to the GeoTIFF file.
+        
+    Returns:
+        tuple: (data, transform, profile) where:
+            - data: numpy array of raster data
+            - transform: Affine transform object
+            - profile: Dictionary of raster metadata
+    """
+    with rasterio.open(filepath) as src:
+        data = src.read()
+        transform = src.transform
+        profile = src.profile
+    return data, transform, profile
+
+
+def is_near(value1, value2, tolerance=0.5):
+    """Check if two values are within a tolerance range."""
+    return abs(value1 - value2) <= tolerance
+
+
+def find_neighboring_files(reference_file, corner_dict):
+    """
+    Find neighboring GeoTIFF files based on corner coordinates.
+    
+    Args:
+        reference_file: Path to the reference GeoTIFF file.
+        corner_dict: Dictionary mapping filenames to corner coordinates.
+        
+    Returns:
+        dict: Dictionary with keys ['right', 'bottom_right', 'bottom', 'bottom_left', 
+             'left', 'top_left', 'top', 'top_right'] containing paths to neighboring files.
+    """
+    neighbors = {
+        'right': None, 'bottom_right': None, 'bottom': None, 'bottom_left': None,
+        'left': None, 'top_left': None, 'top': None, 'top_right': None
+    }
+    reference_basename = os.path.basename(reference_file)
+    if reference_basename not in corner_dict:
+        return neighbors
+
+    reference_min_x, reference_min_y, reference_max_x, reference_max_y = corner_dict[reference_basename]
+    height = reference_max_y - reference_min_y
+    width = reference_max_x - reference_min_x
+
+    for filename, corners in corner_dict.items():
+        if filename == reference_basename:
+            continue
+
+        min_x, min_y, max_x, max_y = corners
+
+        # Check for right neighbor
+        if is_near(min_x, reference_max_x) and is_near(min_y, reference_min_y):
+            neighbors['right'] = os.path.join(os.path.dirname(reference_file), filename)
+
+        # Check for bottom_right neighbor
+        if is_near(min_x, reference_max_x) and is_near(min_y, (reference_min_y - height)):
+            neighbors['bottom_right'] = os.path.join(os.path.dirname(reference_file), filename)
+
+        # Check for bottom neighbor
+        if is_near(min_x, reference_min_x) and is_near(min_y, (reference_min_y - height)):
+            neighbors['bottom'] = os.path.join(os.path.dirname(reference_file), filename)
+
+        # Check for bottom_left neighbor
+        if is_near(min_x, reference_min_x - width) and is_near(min_y, (reference_min_y - height)):
+            neighbors['bottom_left'] = os.path.join(os.path.dirname(reference_file), filename)
+
+        # Check for left neighbor
+        if is_near(min_x, reference_min_x - width) and is_near(min_y, reference_min_y):
+            neighbors['left'] = os.path.join(os.path.dirname(reference_file), filename)
+
+        # Check for top_left neighbor
+        if is_near(min_x, reference_min_x - width) and is_near(min_y, reference_max_y):
+            neighbors['top_left'] = os.path.join(os.path.dirname(reference_file), filename)
+
+        # Check for top neighbor
+        if is_near(min_x, reference_min_x) and is_near(min_y, reference_max_y):
+            neighbors['top'] = os.path.join(os.path.dirname(reference_file), filename)
+
+        # Check for top_right neighbor
+        if is_near(min_x, reference_max_x) and is_near(min_y, reference_max_y):
+            neighbors['top_right'] = os.path.join(os.path.dirname(reference_file), filename)
+
+    return neighbors
+
+
+def create_black_tile_like(reference_data, height, width):
+    """Create a black tile with the same properties as the reference data."""
+    count = reference_data.shape[0]
+    return np.zeros((count, height, width), dtype=reference_data.dtype)
+
+
+def merge_geotiffs(reference_filepath, neighbors, output_filepath):
+    """Merge a reference GeoTIFF with its neighboring tiles into a larger mosaic."""
+    reference_data, _, profile = get_raster_info(reference_filepath)
+    _, reference_height, reference_width = reference_data.shape
+
+    output_width = 2 * reference_width
+    output_height = 2 * reference_height
+
+    _, top_transform, _ = get_raster_info(neighbors.get('top') or reference_filepath)
+
+    new_origin_x = top_transform.c - int(reference_width / 2)
+    new_origin_y = top_transform.f + int(reference_height / 2)
+
+    new_transform = Affine(
+        top_transform.a, top_transform.b, new_origin_x,
+        top_transform.d, top_transform.e, new_origin_y
+    )
+
+    with rasterio.open(
+        output_filepath, 'w', driver=profile['driver'], height=output_height, width=output_width,
+        count=profile['count'], dtype=profile['dtype'], transform=new_transform
+    ) as dst:
+        ref_window_offset_col = reference_width // 2
+        ref_window_offset_row = reference_height // 2
+        dst.write(reference_data, window=Window(
+            col_off=ref_window_offset_col,
+            row_off=ref_window_offset_row,
+            width=reference_width,
+            height=reference_height
+        ))
+
+        tile_layout = {
+            'top_left': (0, 0, (reference_height//2, reference_width//2), 
+                        (slice(reference_height//2, None), slice(reference_width//2, None))),
+            'top': (reference_width//2, 0, (reference_height//2, reference_width), 
+                   (slice(reference_height//2, None), slice(None))),
+            'top_right': (3*reference_width//2, 0, (reference_height//2, reference_width//2), 
+                         (slice(reference_height//2, None), slice(0, reference_width//2))),
+            'left': (0, reference_height//2, (reference_height, reference_width//2), 
+                    (slice(None), slice(reference_width//2, None))),
+            'right': (3*reference_width//2, reference_height//2, (reference_height, reference_width//2), 
+                     (slice(None), slice(0, reference_width//2))),
+            'bottom_left': (0, 3*reference_height//2, (reference_height//2, reference_width//2), 
+                           (slice(0, reference_height//2), slice(reference_width//2, None))),
+            'bottom': (reference_width//2, 3*reference_height//2, (reference_height//2, reference_width), 
+                      (slice(0, reference_height//2), slice(None))),
+            'bottom_right': (3*reference_width//2, 3*reference_height//2, (reference_height//2, reference_width//2), 
+                            (slice(0, reference_height//2), slice(0, reference_width//2))),
+        }
+
+        for direction, (offset_col, offset_row, (h, w), slicing) in tile_layout.items():
+            if neighbors[direction]:
+                neighbor_data, _, _ = get_raster_info(neighbors[direction])
+                neighbor_crop = neighbor_data[:, slicing[0], slicing[1]]
+            else:
+                neighbor_crop = create_black_tile_like(reference_data, h, w)
+
+            dst.write(neighbor_crop, window=Window(offset_col, offset_row, w, h))
+
+
+def get_corner_coordinates(filepath):
+    """Get the corner coordinates of a GeoTIFF file."""
+    with rasterio.open(filepath) as src:
+        transform = src.transform
+        res_x, _, min_x, _, res_y, max_y, _, _, _ = transform
+        width, height = src.width, src.height
+        min_y = max_y + res_y * height
+        max_x = min_x + res_x * width
+    return min_x, min_y, max_x, max_y
+
+
+def save_corner_coordinates(filepaths):
+    """Save corner coordinates for a list of GeoTIFF files."""
+    corners = {}
+    for filepath in filepaths:
+        basename = os.path.basename(filepath)
+        min_x, min_y, max_x, max_y = get_corner_coordinates(filepath)
+        corners[basename] = (min_x, min_y, max_x, max_y)
+    return corners
+
+
+def find_lower_east_file(filepaths, corner_dict):
+    """Find the file with the lowest Y and easternmost X coordinates."""
+    lower_east_file = None
+    lowest_y = float('inf')
+    lowest_x = float('inf')
+
+    for filepath in filepaths:
+        basename = os.path.basename(filepath)
+        min_x, min_y, _, _ = corner_dict[basename]
+        if min_y < lowest_y or (min_y == lowest_y and min_x < lowest_x):
+            lowest_y = min_y
+            lowest_x = min_x
+            lower_east_file = filepath
+
+    return lower_east_file
+
+
+def merge_files_from_folder(folder_path_in, folder_path_out):
+    """Merge all GeoTIFF files in a folder into larger mosaics."""
+    path_parts = folder_path_in.split(os.sep)
+    d_folder = next((part for part in path_parts if part.startswith('D')), '')
+    z_folder = next((part for part in path_parts if part.startswith('Z')), '')
+
+    filepaths = [os.path.join(folder_path_in, f) for f in os.listdir(folder_path_in) 
+                if f.lower().endswith(".tif")]
+
+    corner_dict = save_corner_coordinates(filepaths)
+
+    while filepaths:
+        reference_filepath = find_lower_east_file(filepaths, corner_dict)
+        filepaths.remove(reference_filepath)
+
+        neighbors = find_neighboring_files(reference_filepath, corner_dict)
+        new_basename = f"{d_folder}_{z_folder}_{os.path.basename(reference_filepath)}"
+        output_filepath = os.path.join(folder_path_out, new_basename)
+        merge_geotiffs(reference_filepath, neighbors, output_filepath)
+
+
+def list_folders(path):
+    """List all folders in a directory."""
+    return [f for f in os.listdir(path) if os.path.isdir(os.path.join(path, f))]
+
+
+def process_d_folder(args):
+    """Process a single D folder."""
+    split, modality, d_folder, data_dir, merged_dir = args
+    results = []
+    for z_folder in list_folders(os.path.join(data_dir, split, modality, d_folder)):
+        folder_in = os.path.join(data_dir, split, modality, d_folder, z_folder)
+        folder_out = os.path.join(merged_dir, split, modality)
+        os.makedirs(folder_out, exist_ok=True)
+        merge_files_from_folder(folder_in, folder_out)
+        results.append((folder_in, folder_out))
+    return results
+
+
+if __name__ == "__main__":
+    data_dir = "/mnt/Data/FlairHUB/data_flairhub/output/"
+    merged_dir = "/mnt/Data/FlairHUB/data_flairhub/output/FLAIR1024_optimal"
+    os.makedirs(merged_dir, exist_ok=True)
+
+    for split in ["train", "valid", "test"]:
+        print(f"📂 Split: {split}")
+        for modality in ["img", "msk"]:
+            print(f"🏞️ Modality: {modality}")
+            d_folders = [f for f in list_folders(os.path.join(data_dir, split, modality)) 
+                        if f.startswith("D")]
+
+            args_list = [(split, modality, d, data_dir, merged_dir) for d in d_folders]
+
+            with ProcessPoolExecutor(max_workers=16) as executor:
+                list(tqdm(executor.map(process_d_folder, args_list), total=len(args_list), 
+                     desc="Processing D-folders"))
+
+    print("📝 Done!")
+

main.py

@@ -0,0 +1,89 @@
+"""
+Main entry point for CASWiT training and evaluation.
+
+This script provides a unified interface for training, evaluation, and inference.
+"""
+
+import argparse
+import sys
+import logging
+from pathlib import Path
+
+# Add project root to path
+sys.path.insert(0, str(Path(__file__).parent))
+
+from train.train import main as train_main, load_config
+from train.eval import evaluate_model
+from train.inference import inference_single_image
+from model.CASWiT import CASWiT
+from dataset.definition_dataset import build_transforms
+
+
+def main():
+    """Main entry point."""
+    parser = argparse.ArgumentParser(description="CASWiT: Context-Aware Swin Transformer")
+    parser.add_argument("mode", choices=["train", "eval", "inference"], 
+                       help="Mode: train, eval, or inference")
+    parser.add_argument("--config", type=str, required=True,
+                       help="Path to config YAML file")
+    parser.add_argument("--checkpoint", type=str, default="",
+                       help="Path to model checkpoint (for eval/inference)")
+    parser.add_argument("--image", type=str, default="",
+                       help="Path to input image (for inference)")
+    parser.add_argument("--output", type=str, default="prediction.png",
+                       help="Path to save output (for inference)")
+    parser.add_argument("--split", type=str, default="test", choices=["test", "val"],
+                       help="Dataset split for evaluation")
+    
+    args = parser.parse_args()
+    
+    if args.mode == "train":
+        train_main(args.config)
+    elif args.mode == "eval":
+        if not args.checkpoint:
+            print("Error: --checkpoint required for evaluation")
+            sys.exit(1)
+        cfg = load_config(args.config)
+        evaluate_model(cfg, args.checkpoint, args.split)
+    elif args.mode == "inference":
+        if not args.checkpoint or not args.image:
+            print("Error: --checkpoint and --image required for inference")
+            sys.exit(1)
+        import torch
+        
+        cfg = load_config(args.config)
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        # Validate checkpoint path
+        checkpoint_path_obj = Path(args.checkpoint)
+        if not checkpoint_path_obj.exists() or not checkpoint_path_obj.is_file():
+            print(f"Error: Checkpoint file not found: {args.checkpoint}")
+            sys.exit(1)
+        
+        model = CASWiT(
+            num_head_xa=cfg.cross_attention_heads,
+            num_classes=cfg.num_classes,
+            model_name=cfg.model_name,
+            mlp_ratio=cfg.fusion_mlp_ratio,
+            drop_path=cfg.fusion_drop_path
+        ).to(device)
+        
+        print(f"Loading checkpoint from: {args.checkpoint}")
+        state_dict = torch.load(args.checkpoint, map_location=device)
+        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+        missing, unexpected = model.load_state_dict(state_dict, strict=False)
+        print(f"Successfully loaded checkpoint from: {args.checkpoint}")
+        if len(missing) > 0:
+            print(f"   Missing keys: {len(missing)}")
+        if len(unexpected) > 0:
+            print(f"   Unexpected keys: {len(unexpected)}")
+        if len(missing) == 0 and len(unexpected) == 0:
+            print(f"   Perfect match! All weights loaded successfully.")
+        
+        transform = build_transforms()
+        inference_single_image(model, args.image, device, transform, args.output)
+
+
+if __name__ == "__main__":
+    main()
+

model/CASWiT.py

@@ -0,0 +1,246 @@
+"""
+CASWiT: Context-Aware Swin Transformer for Ultra-High Resolution Semantic Segmentation
+
+This module implements the main CASWiT model architecture with dual-branch
+high-resolution and low-resolution processing with cross-attention fusion.
+"""
+
+import math
+from typing import Dict
+import torch
+import torch.nn as nn
+from transformers import UperNetForSemanticSegmentation
+
+
+class DropPath(nn.Module):
+    """Drop path (stochastic depth) regularization module."""
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        self.drop_prob = float(drop_prob)
+    
+    def forward(self, x):
+        if self.drop_prob == 0.0 or (not self.training):
+            return x
+        keep = 1.0 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        mask = x.new_empty(shape).bernoulli_(keep).div_(keep)
+        return x * mask
+
+
+class CrossFusionBlock(nn.Module):
+    """
+    Cross-attention fusion block that enables HR features to attend to LR features.
+    
+    Implements pre-norm cross-attention (Q=HR, K/V=LR).
+    
+    Args:
+        C_hr: Channel dimension of HR features
+        C_lr: Channel dimension of LR features
+        num_heads: Number of attention heads
+        mlp_ratio: MLP expansion ratio
+        drop: Dropout rate
+        drop_path: Drop path rate
+    """
+    def __init__(self, C_hr: int, C_lr: int, num_heads: int = 8,
+                 mlp_ratio: float = 4.0, drop: float = 0.0, drop_path: float = 0.1):
+        super().__init__()
+
+        self.norm_q = nn.LayerNorm(C_hr)
+        self.norm_kv = nn.LayerNorm(C_lr)
+        self.attn = nn.MultiheadAttention(
+            embed_dim=C_hr, num_heads=num_heads, kdim=C_lr, vdim=C_lr,
+            dropout=drop, batch_first=True
+        )
+
+        hidden = int(C_hr * mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.LayerNorm(C_hr),
+            nn.Linear(C_hr, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, C_hr),
+        )
+
+    def forward(self, x_hr: torch.Tensor, x_lr: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass through cross-attention fusion block.
+        
+        Args:
+            x_hr: HR features [B, C_hr, H_hr, W_hr]
+            x_lr: LR features [B, C_lr, H_lr, W_lr]
+            
+        Returns:
+            Fused HR features [B, C_hr, H_hr, W_hr]
+        """
+        B, C_hr, H_hr, W_hr = x_hr.shape
+        _, C_lr, H_lr, W_lr = x_lr.shape
+
+        # Flatten to sequences
+        q  = x_hr.flatten(2).transpose(1, 2)    # [B, N_hr, C_hr]
+        kv = x_lr.flatten(2).transpose(1, 2)    # [B, N_lr, C_lr]
+
+        # Pre-norm
+        qn  = self.norm_q(q)
+        kvn = self.norm_kv(kv)
+
+        attn_out, _ = self.attn(qn, kvn, kvn)  # [B, N_hr, C_hr]
+
+        # Residual connection + MLP
+        y = q + attn_out
+        y = y + self.mlp(y)
+        
+        return y.transpose(1, 2).view(B, C_hr, H_hr, W_hr)
+
+
+class CASWiT(nn.Module):
+    """
+    CASWiT: Context-Aware Swin Transformer for Ultra-High Resolution Semantic Segmentation.
+    
+    Dual-branch architecture with:
+    - HR branch: Processes high-resolution crops
+    - LR branch: Processes low-resolution context
+    - Cross-attention fusion at each encoder stage
+    
+    Args:
+        num_head_xa: Number of cross-attention heads
+        num_classes: Number of segmentation classes
+        model_name: HuggingFace model identifier for UPerNet-Swin
+        mlp_ratio: MLP expansion ratio in fusion blocks
+        drop_path: Drop path rate
+    """
+    def __init__(self, num_head_xa: int = 1, num_classes: int = 12, 
+                 model_name: str = "openmmlab/upernet-swin-tiny", 
+                 mlp_ratio: float = 4.0, drop_path: float = 0.1):
+        super().__init__()
+        # Load two UPerNet backbones (HR and LR branches)
+        model_hr = UperNetForSemanticSegmentation.from_pretrained(
+            model_name, num_labels=num_classes, ignore_mismatched_sizes=True
+        )
+        model_lr = UperNetForSemanticSegmentation.from_pretrained(
+            model_name, num_labels=num_classes, ignore_mismatched_sizes=True
+        )
+        
+        # Extract HR branch components
+        self.embeddings_hr = model_hr.backbone.embeddings
+        self.encoder_layers_hr = model_hr.backbone.encoder.layers
+        self.hidden_states_norms_hr = model_hr.backbone.hidden_states_norms
+        self.decoder = model_hr.decode_head
+
+        # Extract LR branch components
+        self.embeddings_lr = model_lr.backbone.embeddings
+        self.encoder_layers_lr = model_lr.backbone.encoder.layers
+        self.hidden_states_norms_lr = model_lr.backbone.hidden_states_norms
+        self.decoder_lr = model_lr.decode_head
+
+        # Cross-attention blocks at each stage
+        # Dimensions: tiny:[96, 192, 384, 768] base:[128, 256, 512, 1024] large:[192, 384, 768, 1536]
+        dims_map = {
+            "tiny": [96, 192, 384, 768],
+            "base": [128, 256, 512, 1024],
+            "large": [192, 384, 768, 1536]
+        }
+        # Infer dimensions from model name
+        if "tiny" in model_name.lower():
+            dims = dims_map["tiny"]
+        elif "large" in model_name.lower():
+            dims = dims_map["large"]
+        else:
+            dims = dims_map["base"]  # default to base
+        
+        self.cross_attn_blocks = nn.ModuleList([
+            CrossFusionBlock(dim, dim, num_heads=num_head_xa,
+                           mlp_ratio=mlp_ratio, drop=0.0, drop_path=drop_path) 
+            for dim in dims
+        ])
+
+    def forward(self, x_hr: torch.Tensor, x_lr: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """
+        Forward pass through CASWiT model.
+        
+        Args:
+            x_hr: HR input images [B, 3, H_hr, W_hr]
+            x_lr: LR input images [B, 3, H_lr, W_lr]
+            
+        Returns:
+            Dictionary with 'logits_hr' and 'logits_lr' segmentation logits
+        """
+        B = x_hr.size(0)
+        
+        # Patch embeddings
+        x_hr_seq, _ = self.embeddings_hr(x_hr)
+        x_lr_seq, _ = self.embeddings_lr(x_lr)
+
+        N_hr, C_hr = x_hr_seq.shape[1], x_hr_seq.shape[2]
+        N_lr, C_lr = x_lr_seq.shape[1], x_lr_seq.shape[2]
+        H_hr = W_hr = int(math.sqrt(N_hr))
+        H_lr = W_lr = int(math.sqrt(N_lr))
+        dims_hr = (H_hr, W_hr)
+        dims_lr = (H_lr, W_lr)
+
+        features_hr: Dict[str, torch.Tensor] = {}
+        features_lr: Dict[str, torch.Tensor] = {}
+
+        # Process through encoder stages with cross-attention fusion
+        for idx, (stage_hr, stage_lr, ca) in enumerate(zip(
+            self.encoder_layers_hr, self.encoder_layers_lr, self.cross_attn_blocks
+        )):
+            # HR branch blocks
+            for block in stage_hr.blocks:
+                x_hr_seq = block(x_hr_seq, dims_hr)
+                if isinstance(x_hr_seq, tuple):
+                    x_hr_seq = x_hr_seq[0]
+            
+            # LR branch blocks
+            for block in stage_lr.blocks:
+                x_lr_seq = block(x_lr_seq, dims_lr)
+                if isinstance(x_lr_seq, tuple):
+                    x_lr_seq = x_lr_seq[0]
+
+            # Layer normalization
+            x_hr_seq = self.hidden_states_norms_hr[f"stage{idx+1}"](x_hr_seq)
+            x_lr_seq = self.hidden_states_norms_lr[f"stage{idx+1}"](x_lr_seq)
+
+            H_hr, W_hr = dims_hr
+            H_lr, W_lr = dims_lr
+            C_hr = x_hr_seq.shape[-1]
+            C_lr = x_lr_seq.shape[-1]
+
+            # Reshape to spatial format
+            feat_hr = x_hr_seq.transpose(1, 2).contiguous().view(B, C_hr, H_hr, W_hr)
+            feat_lr = x_lr_seq.transpose(1, 2).contiguous().view(B, C_lr, H_lr, W_lr)
+
+            # Cross-attend HR to LR
+            fused_hr = ca(feat_hr, feat_lr)
+            fused_hr_seq = fused_hr.flatten(2).transpose(1, 2).contiguous()
+
+            # Downsample if stage has it
+            if stage_hr.downsample is not None:
+                fused_hr_seq = stage_hr.downsample(fused_hr_seq, dims_hr)
+                dims_hr = (dims_hr[0] // 2, dims_hr[1] // 2)
+            if stage_lr.downsample is not None:
+                x_lr_seq = stage_lr.downsample(x_lr_seq, dims_lr)
+                dims_lr = (dims_lr[0] // 2, dims_lr[1] // 2)
+
+            features_hr[f"stage{idx+1}"] = fused_hr
+            features_lr[f"stage{idx+1}"] = feat_lr
+            x_hr_seq = fused_hr_seq
+
+        # Decode HR features
+        features_tuple = (
+            features_hr["stage1"],
+            features_hr["stage2"],
+            features_hr["stage3"],
+            features_hr["stage4"],
+        )
+        logits = self.decoder(features_tuple)
+        
+        # Decode LR features (for auxiliary supervision)
+        features_tuple_lr = (
+            features_lr["stage1"],
+            features_lr["stage2"],
+            features_lr["stage3"],
+            features_lr["stage4"],
+        )
+        logits_lr = self.decoder_lr(features_tuple_lr)
+        
+        return {"logits_hr": logits, "logits_lr": logits_lr}
+

model/CASWiT_ssl.py

@@ -0,0 +1,287 @@
+"""
+CASWiT Self-Supervised Learning (SSL) Module
+
+Implements SimMIM-based self-supervised pre-training for CASWiT using
+masked image modeling with dual-branch HR/LR processing.
+"""
+
+import math
+from typing import Optional, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import UperNetForSemanticSegmentation
+
+
+def random_masking_with_tokens(x: torch.Tensor, mask_ratio: float = 0.75, 
+                                mask_token: Optional[torch.Tensor] = None):
+    """
+    Random masking at token level with learned mask token.
+    
+    Args:
+        x: Input tokens [B, N, C]
+        mask_ratio: Ratio of tokens to mask
+        mask_token: Learnable mask token
+        
+    Returns:
+        x_masked: Masked tokens [B, N, C]
+        mask: Binary mask [B, N] where 0=visible, 1=masked
+        ids_restore: Indices to restore original order
+    """
+    B, N, C = x.shape
+    len_keep = int(N * (1 - mask_ratio))
+
+    noise = torch.rand(B, N, device=x.device)
+    ids_shuffle = torch.argsort(noise, dim=1)
+    ids_restore = torch.argsort(ids_shuffle, dim=1)
+    ids_keep = ids_shuffle[:, :len_keep]
+
+    x_keep = torch.gather(x, 1, ids_keep.unsqueeze(-1).expand(-1, -1, C))
+
+    if mask_token is None:
+        mask_token = torch.zeros((1, C), device=x.device)
+    m_tok = mask_token.view(1, 1, C).expand(B, N - len_keep, C)
+
+    x_cat = torch.cat([x_keep, m_tok], dim=1)
+    x_masked = torch.gather(x_cat, 1, ids_restore.unsqueeze(-1).expand(-1, -1, C))
+
+    mask = torch.ones(B, N, device=x.device)
+    mask[:, :len_keep] = 0
+    mask = torch.gather(mask, 1, ids_restore)
+    return x_masked, mask, ids_restore
+
+
+def center_masking_with_tokens(x: torch.Tensor, mask_token: Optional[torch.Tensor] = None, 
+                                mask_ratio: float = 0.5):
+    """
+    Deterministic centered square mask.
+    
+    Args:
+        x: Input tokens [B, N, C]
+        mask_token: Learnable mask token
+        mask_ratio: Ratio of tokens to mask
+        
+    Returns:
+        x_masked: Masked tokens [B, N, C]
+        mask: Binary mask [B, N]
+        ids_restore: Indices to restore original order
+    """
+    B, N, C = x.shape
+    H = W = int(N**0.5)
+    assert H * W == N, "N must be a perfect square"
+    L = int(round(H * (mask_ratio ** 0.5)))
+    start = (H - L) // 2
+    end = start + L
+
+    mask_2d = torch.zeros(H, W, device=x.device, dtype=torch.bool)
+    mask_2d[start:end, start:end] = True
+    mask = mask_2d.view(1, -1).expand(B, -1)  # (B,N)
+
+    if mask_token is None:
+        mask_token = torch.zeros(C, device=x.device)
+    mask_token = mask_token.view(-1)
+
+    x_masked = x * (~mask).unsqueeze(-1) + mask.unsqueeze(-1) * mask_token.view(1, 1, C)
+    ids_restore = torch.arange(N, device=x.device).unsqueeze(0).expand(B, N)
+    return x_masked, mask.to(x_masked.dtype), ids_restore
+
+
+class CrossAttentionBlock(nn.Module):
+    """Simplified cross-attention block for SSL."""
+    def __init__(self, C_hr, C_lr, num_heads=8, dropout=0.0):
+        super().__init__()
+        self.cross_attn = nn.MultiheadAttention(
+            embed_dim=C_hr, num_heads=num_heads, kdim=C_lr, vdim=C_lr, 
+            dropout=dropout, batch_first=True
+        )
+        self.norm = nn.LayerNorm(C_hr)
+        self.mlp = nn.Sequential(
+            nn.LayerNorm(C_hr),
+            nn.Linear(C_hr, C_hr * 4),
+            nn.GELU(),
+            nn.Linear(C_hr * 4, C_hr),
+        )
+
+    def forward(self, x_hr, x_lr):
+        B, C_hr, H_hr, W_hr = x_hr.shape
+        _, C_lr, H_lr, W_lr = x_lr.shape
+        q = x_hr.flatten(2).transpose(1, 2)  # (B,N_hr,C_hr)
+        kv = x_lr.flatten(2).transpose(1, 2)  # (B,N_lr,C_lr)
+        attn_out, _ = self.cross_attn(q, kv, kv)
+        y = self.norm(q + attn_out)
+        y = y + self.mlp(y)
+        return y.transpose(1, 2).view(B, C_hr, H_hr, W_hr)
+
+
+class CASWiT_SSL(nn.Module):
+    """
+    CASWiT Self-Supervised Learning model using SimMIM.
+    
+    Encoder: Dual Swin backbones with cross-attention blocks
+    Decoder: Conv1x1 + PixelShuffle for reconstruction
+    Masking: HR random masking, LR center masking
+    
+    Args:
+        model_name: HuggingFace model identifier
+        mask_ratio_hr: Masking ratio for HR branch
+        mask_ratio_lr: Masking ratio for LR branch
+        patch_size: Patch size for masking
+        encoder_stride: Encoder stride for decoder
+        xa_heads: Number of cross-attention heads per stage
+    """
+    def __init__(self, model_name: str = "openmmlab/upernet-swin-base", 
+                 mask_ratio_hr: float = 0.75, mask_ratio_lr: float = 0.5,
+                 patch_size: int = 4, encoder_stride: int = 32,
+                 xa_heads: Tuple[int, int, int, int] = (8, 8, 8, 8)):
+        super().__init__()
+        self.mask_ratio_hr = mask_ratio_hr
+        self.mask_ratio_lr = mask_ratio_lr
+        self.patch_size = patch_size
+        self.encoder_stride = encoder_stride
+
+        # Load two UPerNet (Swin) backbones
+        model_hr = UperNetForSemanticSegmentation.from_pretrained(
+            model_name, ignore_mismatched_sizes=True
+        )
+        model_lr = UperNetForSemanticSegmentation.from_pretrained(
+            model_name, ignore_mismatched_sizes=True
+        )
+
+        self.embeddings_hr = model_hr.backbone.embeddings
+        self.encoder_layers_hr = model_hr.backbone.encoder.layers
+        self.hidden_states_norms_hr = model_hr.backbone.hidden_states_norms
+
+        self.embeddings_lr = model_lr.backbone.embeddings
+        self.encoder_layers_lr = model_lr.backbone.encoder.layers
+        self.hidden_states_norms_lr = model_lr.backbone.hidden_states_norms
+
+        # Cross-attention blocks with explicit Swin-Base dims
+        dims = [128, 256, 512, 1024]
+        self.cross_attn_blocks = nn.ModuleList([
+            CrossAttentionBlock(d, d, num_heads=h) for d, h in zip(dims, xa_heads)
+        ])
+
+        # Learnable mask tokens
+        self.mask_token_hr = nn.Parameter(torch.zeros(1, dims[0]))
+        self.mask_token_lr = nn.Parameter(torch.zeros(1, dims[0]))
+
+        # SimMIM decoder: Conv1×1 → PixelShuffle(stride)
+        self.decoder_conv = None  # lazy init after we know C_last
+        self.decoder_shuffle = nn.PixelShuffle(self.encoder_stride)
+
+        # Store masks for visualization
+        self.last_mask_hr = None
+        self.last_mask_lr = None
+
+    def _encode(self, x_hr: torch.Tensor, x_lr: torch.Tensor):
+        """Encode with masking and return reconstruction targets."""
+        B, C, H, W = x_hr.shape
+        target_img = x_hr
+        target_lr = x_lr
+
+        # Patch embeddings
+        x_hr_seq, _ = self.embeddings_hr(x_hr)  # (B, N_hr, C1)
+        x_lr_seq, _ = self.embeddings_lr(x_lr)  # (B, N_lr, C1)
+
+        # Masking
+        x_hr_seq, mask_hr, _ = random_masking_with_tokens(
+            x_hr_seq, self.mask_ratio_hr, self.mask_token_hr
+        )
+        x_lr_seq, mask_lr, _ = center_masking_with_tokens(
+            x_lr_seq, self.mask_token_lr, mask_ratio=self.mask_ratio_lr
+        )
+
+        # Initial spatial dims
+        H_hr = W_hr = int(math.sqrt(x_hr_seq.shape[1]))
+        H_lr = W_lr = int(math.sqrt(x_lr_seq.shape[1]))
+        dims_hr = (H_hr, W_hr)
+        dims_lr = (H_lr, W_lr)
+
+        # Walk encoder stages with cross attention at each stage
+        for idx, (stage_hr, stage_lr, ca) in enumerate(zip(
+            self.encoder_layers_hr, self.encoder_layers_lr, self.cross_attn_blocks
+        )):
+            # HR blocks
+            for block in stage_hr.blocks:
+                x_hr_seq = block(x_hr_seq, dims_hr)
+                if isinstance(x_hr_seq, tuple):
+                    x_hr_seq = x_hr_seq[0]
+            # LR blocks
+            for block in stage_lr.blocks:
+                x_lr_seq = block(x_lr_seq, dims_lr)
+                if isinstance(x_lr_seq, tuple):
+                    x_lr_seq = x_lr_seq[0]
+
+            # Norms
+            x_hr_seq = self.hidden_states_norms_hr[f"stage{idx+1}"](x_hr_seq)
+            x_lr_seq = self.hidden_states_norms_lr[f"stage{idx+1}"](x_lr_seq)
+
+            # Maps
+            B_, N_hr_, C_hr_ = x_hr_seq.shape
+            B_, N_lr_, C_lr_ = x_lr_seq.shape
+            Hh, Wh = dims_hr
+            Hl, Wl = dims_lr
+            feat_hr = x_hr_seq.transpose(1, 2).contiguous().view(B_, C_hr_, Hh, Wh)
+            feat_lr = x_lr_seq.transpose(1, 2).contiguous().view(B_, C_lr_, Hl, Wl)
+
+            # Cross-fuse HR <- LR
+            fused_hr = ca(feat_hr, feat_lr)
+            x_hr_seq = fused_hr.flatten(2).transpose(1, 2).contiguous()
+
+            # Downsample to next stage
+            if stage_hr.downsample is not None:
+                x_hr_seq = stage_hr.downsample(x_hr_seq, dims_hr)
+                dims_hr = (dims_hr[0] // 2, dims_hr[1] // 2)
+            if stage_lr.downsample is not None:
+                x_lr_seq = stage_lr.downsample(x_lr_seq, dims_lr)
+                dims_lr = (dims_lr[0] // 2, dims_lr[1] // 2)
+
+        # Last-stage feature map z (B, C_last, H/stride, W/stride)
+        Hs, Ws = dims_hr
+        C_last = x_hr_seq.shape[-1]
+        z = x_hr_seq.transpose(1, 2).contiguous().view(B, C_last, Hs, Ws)
+
+        # Lazy init decoder conv
+        if self.decoder_conv is None:
+            self.decoder_conv = nn.Conv2d(
+                C_last, (self.encoder_stride ** 2) * 3, kernel_size=1
+            ).to(z.device)
+
+        # Reconstruction
+        x_rec = self.decoder_shuffle(self.decoder_conv(z))  # (B,3,H,W)
+
+        # Convert patch masks to pixel masks
+        Mh = int(math.sqrt(mask_hr.shape[1]))
+        mask_patch_hr = mask_hr.view(B, Mh, Mh)
+        mask_pix_hr = mask_patch_hr.repeat_interleave(
+            self.patch_size, 1
+        ).repeat_interleave(self.patch_size, 2).unsqueeze(1).contiguous()
+
+        Ml = int(math.sqrt(mask_lr.shape[1]))
+        mask_patch_lr = mask_lr.view(B, Ml, Ml)
+        mask_pix_lr = mask_patch_lr.repeat_interleave(
+            self.patch_size, 1
+        ).repeat_interleave(self.patch_size, 2).unsqueeze(1).contiguous()
+
+        self.last_mask_hr = mask_patch_hr
+        self.last_mask_lr = mask_patch_lr
+
+        return x_rec, target_img, mask_pix_hr, target_lr, mask_pix_lr
+
+    def forward(self, x_hr: torch.Tensor, x_lr: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass for SSL training.
+        
+        Returns reconstruction loss on masked pixels only.
+        """
+        x_rec, target_img, mask_pix, _, _ = self._encode(x_hr, x_lr)
+        loss_recon = F.l1_loss(target_img, x_rec, reduction='none')
+        loss = (loss_recon * mask_pix).sum() / (mask_pix.sum() + 1e-6) / target_img.shape[1]
+        return loss
+
+    @torch.no_grad()
+    def forward_outputs(self, x_hr: torch.Tensor, x_lr: torch.Tensor):
+        """Forward pass returning all outputs for visualization."""
+        x_rec, target_img, mask_pix_hr, target_lr, mask_pix_lr = self._encode(x_hr, x_lr)
+        return x_rec, target_img, mask_pix_hr, target_lr, mask_pix_lr
+

model/__init__.py

@@ -0,0 +1,9 @@
+"""
+CASWiT model implementations.
+"""
+
+from model.CASWiT import CASWiT, CrossFusionBlock
+from model.CASWiT_ssl import CASWiT_SSL
+
+__all__ = ['CASWiT', 'CrossFusionBlock', 'CASWiT_SSL']
+

requirements.txt

@@ -0,0 +1,27 @@
+# Core dependencies
+torch>=2.0.0
+torchvision>=0.15.0
+transformers>=4.30.0
+timm>=0.9.0
+
+# Data processing
+numpy>=1.21.0
+pillow>=9.0.0
+tifffile>=2023.0.0
+opencv-python>=4.5.0
+pandas>=1.3.0
+rasterio>=1.3.0
+affine>=2.3.0
+
+# Training utilities
+tqdm>=4.64.0
+PyYAML>=6.0
+scikit-learn>=1.0.0
+
+# Logging and visualization
+wandb>=0.15.0
+matplotlib>=3.5.0
+
+# Optional: for distributed training
+accelerate>=0.20.0
+

train/__init__.py

@@ -0,0 +1,8 @@
+"""
+Training, evaluation, and inference scripts.
+"""
+
+from train.train import TrainConfig, load_config
+
+__all__ = ['TrainConfig', 'load_config']
+

train/eval.py

@@ -0,0 +1,153 @@
+"""
+Evaluation script for CASWiT model.
+
+Evaluates a trained model on test/validation sets and computes metrics.
+"""
+
+import sys
+import yaml
+import logging
+from pathlib import Path
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torch.cuda.amp import autocast
+from tqdm import tqdm
+
+# Add project root to Python path
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from model.CASWiT import CASWiT
+from dataset.definition_dataset import SemanticSegmentationDatasetFusion, URURHRLRDataset, build_transforms
+from utils.metrics import compute_metrics_from_confusion
+from train.train import load_config, TrainConfig
+
+
+def evaluate_model(cfg: TrainConfig, checkpoint_path: str, split: str = "test"):
+    """
+    Evaluate model on specified split.
+    
+    Args:
+        cfg: Training configuration
+        checkpoint_path: Path to model checkpoint
+        split: Dataset split to evaluate ('test' or 'val')
+    """
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Validate checkpoint path
+    checkpoint_path_obj = Path(checkpoint_path)
+    if not checkpoint_path_obj.exists() or not checkpoint_path_obj.is_file():
+        raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+    
+    # Load model
+    model = CASWiT(
+        num_head_xa=cfg.cross_attention_heads,
+        num_classes=cfg.num_classes,
+        model_name=cfg.model_name,
+        mlp_ratio=cfg.fusion_mlp_ratio,
+        drop_path=cfg.fusion_drop_path
+    ).to(device)
+    
+    # Load checkpoint
+    print(f"Loading checkpoint from: {checkpoint_path}")
+    state_dict = torch.load(checkpoint_path, map_location=device)
+    state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    print(f"Successfully loaded checkpoint from: {checkpoint_path}")
+    if len(missing) > 0:
+        print(f"   Missing keys: {len(missing)}")
+    if len(unexpected) > 0:
+        print(f"   Unexpected keys: {len(unexpected)}")
+    if len(missing) == 0 and len(unexpected) == 0:
+        print(f"   Perfect match! All weights loaded successfully.")
+    model.eval()
+    dataset_name = cfg.dataset_name
+    
+    # Create dataset with sliding window for test, simple crop for val
+    t = build_transforms()
+    base = Path(cfg.data_path)
+    if dataset_name == "URUR":
+        # Use sliding window without tiling for URUR (full image coverage)
+        ds = URURHRLRDataset(
+            image_dir=base / cfg.test_img_subdir,
+            mask_dir=base / cfg.test_msk_subdir,
+            num_classes=cfg.num_classes,
+            mode="test",
+            ignore_index=cfg.ignore_index,
+            transform=t
+        )
+    else:
+        # Use simple center crop with FLAIRHUB
+        ds = SemanticSegmentationDatasetFusion(base / cfg.test_img_subdir, base / cfg.test_msk_subdir, transform=t)
+    
+    dl = DataLoader(ds, batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.num_workers, pin_memory=True)
+    
+    # Evaluate
+    criterion = torch.nn.CrossEntropyLoss(ignore_index=cfg.ignore_index)
+    running_loss = 0.0
+    full_confmat = torch.zeros((cfg.num_classes, cfg.num_classes), dtype=torch.long, device=device)
+    
+    with torch.inference_mode():
+        for batch in tqdm(dl, desc=f"Evaluating {split}"):
+            # Handle datasets that return meta dict (URURHRLRDataset returns 5 values)
+            if len(batch) == 5:
+                images_hr, masks_hr, images_lr, masks_lr, _ = batch
+            else:
+                images_hr, masks_hr, images_lr, masks_lr = batch
+            images_hr = images_hr.to(device, non_blocking=True)
+            masks_hr = masks_hr.to(device, non_blocking=True)
+            images_lr = images_lr.to(device, non_blocking=True)
+            masks_lr = masks_lr.to(device, non_blocking=True)
+            
+            with autocast():
+                out = model(images_hr, images_lr)
+                logits_hr = out["logits_hr"]
+                logits_hr = F.interpolate(logits_hr, size=masks_hr.shape[-2:], mode="bilinear", align_corners=False)
+                loss = criterion(logits_hr, masks_hr)
+            
+            running_loss += float(loss.item())
+            
+            preds = torch.argmax(logits_hr, dim=1)
+            valid = (masks_hr >= 0) & (masks_hr < cfg.num_classes)
+            t = masks_hr[valid]
+            p = preds[valid]
+            
+            cm = torch.bincount(
+                (t * cfg.num_classes + p).view(-1),
+                minlength=cfg.num_classes * cfg.num_classes
+            ).reshape(cfg.num_classes, cfg.num_classes)
+            full_confmat += cm
+    
+    avg_loss = running_loss / len(dl)
+    confmat_np = full_confmat.cpu().numpy()
+    metrics = compute_metrics_from_confusion(confmat_np)
+    
+    print(f"\n{split.upper()} Results:")
+    print(f"  Loss: {avg_loss:.4f}")
+    print(f"  mIoU: {metrics['mIoU']:.4f}")
+    print(f"  mF1: {metrics['mF1']:.4f}")
+    print(f"  Per-class IoU: {metrics['IoUs']}")
+    
+    return metrics
+
+
+def main():
+    """Main evaluation function."""
+    import sys
+    if len(sys.argv) < 3:
+        print("Usage: python eval.py <config_path> <checkpoint_path> [split]")
+        sys.exit(1)
+    
+    cfg_path = sys.argv[1]
+    checkpoint_path = sys.argv[2]
+    split = sys.argv[3] if len(sys.argv) > 3 else "test"
+    
+    logging.basicConfig(level=logging.INFO)
+    cfg = load_config(cfg_path)
+    evaluate_model(cfg, checkpoint_path, split)
+
+
+if __name__ == "__main__":
+    main()
+

train/inference.py

@@ -0,0 +1,118 @@
+"""
+Inference script for CASWiT model.
+
+Performs inference on images and saves predictions.
+"""
+
+import sys
+import yaml
+import logging
+from pathlib import Path
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from tifffile import imread
+from torchvision import transforms
+
+# Add project root to Python path
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from model.CASWiT import CASWiT
+from dataset.definition_dataset import build_transforms
+from train.train import load_config
+
+
+def inference_single_image(model, image_path: str, device, transform, output_path: str = None):
+    """
+    Run inference on a single image.
+    
+    Args:
+        model: Trained CASWiT model
+        image_path: Path to input image
+        device: Device to run inference on
+        transform: Image transform
+        output_path: Path to save prediction (optional)
+    """
+    # Load and preprocess image
+    image = imread(str(image_path))[:, :, :3]
+    
+    # Create HR and LR versions
+    crop_x, crop_y = 256, 256
+    image_hr = image[crop_x:crop_x + 512, crop_y:crop_y + 512]
+    image_lr = image[::2, ::2, :]
+    
+    # Transform
+    img_hr_tensor = transform(Image.fromarray(image_hr)).unsqueeze(0).to(device)
+    img_lr_tensor = transform(Image.fromarray(image_lr)).unsqueeze(0).to(device)
+    
+    # Inference
+    model.eval()
+    with torch.no_grad():
+        out = model(img_hr_tensor, img_lr_tensor)
+        logits = out["logits_hr"]
+        pred = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy()
+    
+    if output_path:
+        # Save prediction as image
+        pred_img = Image.fromarray(pred.astype(np.uint8))
+        pred_img.save(output_path)
+        print(f"Prediction saved to {output_path}")
+    
+    return pred
+
+
+def main():
+    """Main inference function."""
+    import sys
+    if len(sys.argv) < 4:
+        print("Usage: python inference.py <config_path> <checkpoint_path> <image_path> [output_path]")
+        sys.exit(1)
+    
+    cfg_path = sys.argv[1]
+    checkpoint_path = sys.argv[2]
+    image_path = sys.argv[3]
+    output_path = sys.argv[4] if len(sys.argv) > 4 else "prediction.png"
+    
+    logging.basicConfig(level=logging.INFO)
+    cfg = load_config(cfg_path)
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Validate checkpoint path
+    checkpoint_path_obj = Path(checkpoint_path)
+    if not checkpoint_path_obj.exists() or not checkpoint_path_obj.is_file():
+        raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+    
+    # Load model
+    model = CASWiT(
+        num_head_xa=cfg.cross_attention_heads,
+        num_classes=cfg.num_classes,
+        model_name=cfg.model_name,
+        mlp_ratio=cfg.fusion_mlp_ratio,
+        drop_path=cfg.fusion_drop_path
+    ).to(device)
+    
+    # Load checkpoint
+    print(f"Loading checkpoint from: {checkpoint_path}")
+    state_dict = torch.load(checkpoint_path, map_location=device)
+    state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    print(f"Successfully loaded checkpoint from: {checkpoint_path}")
+    if len(missing) > 0:
+        print(f"   Missing keys: {len(missing)}")
+    if len(unexpected) > 0:
+        print(f"   Unexpected keys: {len(unexpected)}")
+    if len(missing) == 0 and len(unexpected) == 0:
+        print(f"   Perfect match! All weights loaded successfully.")
+    
+    # Run inference
+    transform = build_transforms()
+    pred = inference_single_image(model, image_path, device, transform, output_path)
+    print(f"Inference complete. Prediction shape: {pred.shape}")
+
+
+if __name__ == "__main__":
+    main()
+

train/train.py

@@ -0,0 +1,534 @@
+"""
+Training script for CASWiT model.
+
+Supports distributed training with DDP, mixed precision, and WandB logging.
+"""
+
+import os
+import sys
+import yaml
+import logging
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Tuple
+
+# Add project root to Python path
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.distributed as dist
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from torch.cuda.amp import GradScaler, autocast
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from tqdm import tqdm
+
+from model.CASWiT import CASWiT
+from dataset.definition_dataset import SemanticSegmentationDatasetFusion, URURHRLRDataset, build_transforms
+from utils.metrics import compute_metrics_from_confusion
+from utils.logging import setup_wandb_logging, log_metrics
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+
+
+@dataclass
+class TrainConfig:
+    """Training configuration dataclass."""
+    # Paths
+    data_path: str
+    dataset_name: str
+    train_img_subdir: str
+    train_msk_subdir: str
+    val_img_subdir: str
+    val_msk_subdir: str
+    test_img_subdir: str
+    test_msk_subdir: str
+    save_dir: str = "weights"
+    pretrained_path: str = ""
+
+    # Model
+    model_name: str = "openmmlab/upernet-swin-tiny"
+    num_classes: int = 15
+    cross_attention_heads: int = 1
+    ignore_index: int = 255
+    fusion_mlp_ratio: float = 4.0
+    fusion_drop_path: float = 0.1
+    lr_supervision_weight: float = 0.5
+
+    # Training
+    batch_size: int = 4
+    num_workers: int = 4
+    num_epochs: int = 50
+    learning_rate: float = 1e-4
+    amp: bool = True
+    seed: int = 1337
+    eta_min: float = 1e-6
+
+    # Logging
+    use_wandb: bool = True
+    wandb_project: str = "your_project"
+    wandb_entity: str = "your_entity"
+    wandb_run_name: str = "hrlr_fusion"
+
+    # Misc
+    print_device: bool = True
+
+
+def load_config(cfg_path: str) -> TrainConfig:
+    """Load configuration from YAML file."""
+    with open(cfg_path, "r") as f:
+        raw = yaml.safe_load(f)
+
+    training = raw.get("training", {})
+    model = raw.get("model", {})
+    wandb_cfg = raw.get("wandb", {})
+    paths = raw.get("paths", {})
+
+    return TrainConfig(
+        data_path=paths.get("data_path", ""),
+        dataset_name=paths.get("dataset_name", ""),
+        train_img_subdir=paths.get("train_img_subdir", "train/img"),
+        train_msk_subdir=paths.get("train_msk_subdir", "train/msk"),
+        val_img_subdir=paths.get("val_img_subdir", "val/img"),
+        val_msk_subdir=paths.get("val_msk_subdir", "val/msk"),
+        test_img_subdir=paths.get("test_img_subdir", "test/img"),
+        test_msk_subdir=paths.get("test_msk_subdir", "test/msk"),
+        save_dir=paths.get("save_dir", "weights"),
+        pretrained_path=paths.get("pretrained_path", ""),
+        model_name=model.get("model_name", "openmmlab/upernet-swin-tiny"),
+        num_classes=int(model.get("num_classes", 12)),
+        cross_attention_heads=int(model.get("cross_attention_heads", 1)),
+        ignore_index=int(model.get("ignore_index", 255)),
+        fusion_mlp_ratio=float(model.get("fusion_mlp_ratio", 4.0)),
+        fusion_drop_path=float(model.get("fusion_drop_path", 0.1)),
+        lr_supervision_weight=float(training.get("lr_supervision_weight", 0.5)),
+        batch_size=int(training.get("batch_size", 8)),
+        num_workers=int(training.get("num_workers", 4)),
+        num_epochs=int(training.get("num_epochs", 50)),
+        learning_rate=float(training.get("learning_rate", 1e-4)),
+        amp=bool(training.get("amp", True)),
+        seed=int(training.get("seed", 1337)),
+        eta_min=float(training.get("eta_min", 1e-6)),
+        use_wandb=bool(wandb_cfg.get("use_wandb", True)),
+        wandb_project=wandb_cfg.get("project", "your_project"),
+        wandb_entity=wandb_cfg.get("entity", "your_entity"),
+        wandb_run_name=wandb_cfg.get("run_name", "hrlr_fusion"),
+        print_device=bool(raw.get("print_device", True)),
+    )
+
+
+def set_seed(seed: int):
+    """Set random seeds for reproducibility."""
+    import random
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def is_distributed() -> bool:
+    """Check if running in distributed mode."""
+    return int(os.environ.get("WORLD_SIZE", "1")) > 1
+
+
+def get_rank() -> int:
+    """Get current process rank."""
+    return int(os.environ.get("RANK", "0"))
+
+
+def get_local_rank() -> int:
+    """Get local rank."""
+    return int(os.environ.get("LOCAL_RANK", os.environ.get("RANK", "0")))
+
+
+def is_main_process() -> bool:
+    """Check if this is the main process."""
+    return get_rank() == 0
+
+
+def setup_distributed():
+    """Setup distributed training."""
+    if is_distributed() and not dist.is_initialized():
+        dist.init_process_group(backend="nccl")
+        local_rank = get_local_rank()
+        torch.cuda.set_device(local_rank)
+
+
+def cleanup_distributed():
+    """Cleanup distributed training."""
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+def make_dataloaders(cfg: TrainConfig) -> Tuple[DataLoader, DataLoader, DataLoader, 
+                                                  DistributedSampler, DistributedSampler, DistributedSampler]:
+    """Create data loaders for train, val, and test splits."""
+    t = build_transforms()
+    base = Path(cfg.data_path)
+    dataset_name = cfg.dataset_name
+    if dataset_name == "URUR":
+        ds_train = URURHRLRDataset(image_dir=base / cfg.train_img_subdir, mask_dir=base / cfg.train_msk_subdir,
+                             num_classes=cfg.num_classes, mode="train", ignore_index=cfg.ignore_index, 
+                             transform=t)
+        ds_val = URURHRLRDataset(image_dir=base / cfg.val_img_subdir, mask_dir=base / cfg.val_msk_subdir,
+                             num_classes=cfg.num_classes, mode="val", ignore_index=cfg.ignore_index, 
+                             transform=t)
+        ds_test = URURHRLRDataset(image_dir=base / cfg.test_img_subdir, mask_dir=base / cfg.test_msk_subdir,
+                             num_classes=cfg.num_classes, mode="test", ignore_index=cfg.ignore_index, 
+                             transform=t)
+    else: # FLAIRHUB
+        ds_train = SemanticSegmentationDatasetFusion(base / cfg.train_img_subdir, base / cfg.train_msk_subdir, transform=t)
+        ds_val = SemanticSegmentationDatasetFusion(base / cfg.val_img_subdir, base / cfg.val_msk_subdir, transform=t)
+        ds_test = SemanticSegmentationDatasetFusion(base / cfg.test_img_subdir, base / cfg.test_msk_subdir, transform=t)
+
+    if is_distributed():
+        train_sampler = DistributedSampler(ds_train, shuffle=True)
+        val_sampler = DistributedSampler(ds_val, shuffle=False)
+        test_sampler = DistributedSampler(ds_test, shuffle=False)
+        shuffle_train = False
+        shuffle_eval = False
+    else:
+        train_sampler = None
+        val_sampler = None
+        test_sampler = None
+        shuffle_train = True
+        shuffle_eval = False
+
+    dl_train = DataLoader(ds_train, batch_size=cfg.batch_size, sampler=train_sampler,
+                          shuffle=shuffle_train, num_workers=cfg.num_workers, drop_last=True,
+                          pin_memory=True)
+    dl_val = DataLoader(ds_val, batch_size=cfg.batch_size, sampler=val_sampler,
+                        shuffle=shuffle_eval, num_workers=cfg.num_workers, pin_memory=True)
+    dl_test = DataLoader(ds_test, batch_size=cfg.batch_size, sampler=test_sampler,
+                         shuffle=shuffle_eval, num_workers=cfg.num_workers, pin_memory=True)
+
+    return dl_train, dl_val, dl_test, train_sampler, val_sampler, test_sampler
+
+
+def maybe_load_pretrained(model: nn.Module, path: Path):
+    """Load pretrained weights if available."""
+    # Check if path is empty, None, or not a valid file
+    path_str = str(path) if path else ""
+    if not path or path_str.strip() == "" or not path.exists() or not path.is_file():
+        if path and path_str.strip() != "":
+            print(f"WARNING: Pretrained weights not found at {path}. Skipping load.")
+        return
+    
+    # Load the checkpoint
+    print(f"Loading pretrained weights from: {path}")
+    try:
+        state_dict = torch.load(str(path), map_location="cpu")
+        # Remove DistributedDataParallel/DataParallel prefix if any
+        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+        missing, unexpected = model.load_state_dict(state_dict, strict=False)
+        
+        # Log successful loading with details
+        print(f"Successfully loaded pretrained weights from: {path}")
+        if len(missing) > 0:
+            print(f"   Missing keys: {len(missing)} (this is normal if architecture changed)")
+        if len(unexpected) > 0:
+            print(f"   Unexpected keys: {len(unexpected)} (these were ignored)")
+        if len(missing) == 0 and len(unexpected) == 0:
+            print(f"   Perfect match! All weights loaded successfully.")
+    except Exception as e:
+        print(f"ERROR: Failed to load pretrained weights from {path}: {e}")
+        raise
+
+
+def train_one_epoch(model, dl, device, criterion, optimizer, scaler: GradScaler, amp: bool, lr_supervision_weight: float):
+    """Train for one epoch."""
+    model.train()
+    running = 0.0
+    with tqdm(dl, desc="Training", unit="batch", disable=not is_main_process()) as bar:
+        for batch in bar:
+            # Handle datasets that return meta dict (e.g., URURHRLRDataset)
+            if len(batch) == 5:
+                images_hr, masks_hr, images_lr, masks_lr, _ = batch
+            else:
+                images_hr, masks_hr, images_lr, masks_lr = batch
+            images_hr = images_hr.to(device, non_blocking=True)
+            masks_hr = masks_hr.to(device, non_blocking=True)
+            images_lr = images_lr.to(device, non_blocking=True)
+            masks_lr = masks_lr.to(device, non_blocking=True)
+
+            optimizer.zero_grad(set_to_none=True)
+            if amp:
+                with autocast():
+                    out = model(images_hr, images_lr)
+                    logits_hr = out["logits_hr"]
+                    logits_lr = out["logits_lr"]
+                    logits_hr = F.interpolate(logits_hr, size=masks_hr.shape[-2:], mode="bilinear", align_corners=False)
+                    logits_lr = F.interpolate(logits_lr, size=masks_lr.shape[-2:], mode="bilinear", align_corners=False)
+                    loss_hr = criterion(logits_hr, masks_hr)
+                    loss_lr = criterion(logits_lr, masks_lr)
+                    loss = loss_hr + lr_supervision_weight * loss_lr
+                scaler.scale(loss).backward()
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                out = model(images_hr, images_lr)
+                logits_hr = out["logits_hr"]
+                logits_lr = out["logits_lr"]
+                logits_hr = F.interpolate(logits_hr, size=masks_hr.shape[-2:], mode="bilinear", align_corners=False)
+                logits_lr = F.interpolate(logits_lr, size=masks_lr.shape[-2:], mode="bilinear", align_corners=False)
+                loss_hr = criterion(logits_hr, masks_hr)
+                loss_lr = criterion(logits_lr, masks_lr)
+                loss = loss_hr + lr_supervision_weight * loss_lr
+                loss.backward()
+                optimizer.step()
+
+            running += float(loss.item())
+            bar.set_postfix(loss=float(loss.item()))
+    return running / max(1, len(dl))
+
+
+def evaluate(model, dl, device, criterion, num_classes: int, lr_supervision_weight: float, phase_name: str = "Validation"):
+    """Evaluate model on dataset."""
+    model.eval()
+    running = 0.0
+    # Confusion matrix on GPU
+    full_confmat = torch.zeros((num_classes, num_classes), dtype=torch.long, device=device)
+
+    is_main = (not is_distributed()) or is_main_process()
+    iterator = tqdm(dl, desc=phase_name, unit="batch", disable=not is_main) if is_main else dl
+
+    with torch.inference_mode():
+        for batch in iterator:
+            # Handle datasets that return meta dict (e.g., URURHRLRDataset)
+            if len(batch) == 5:
+                images_hr, masks_hr, images_lr, masks_lr, _ = batch
+            else:
+                images_hr, masks_hr, images_lr, masks_lr = batch
+            images_hr = images_hr.to(device, non_blocking=True)
+            masks_hr = masks_hr.to(device, non_blocking=True)
+            images_lr = images_lr.to(device, non_blocking=True)
+            masks_lr = masks_lr.to(device, non_blocking=True)
+
+            with autocast():
+                out = model(images_hr, images_lr)
+                logits_hr = out["logits_hr"]
+                logits_lr = out["logits_lr"]
+                logits_hr = F.interpolate(logits_hr, size=masks_hr.shape[-2:], mode="bilinear", align_corners=False)
+                logits_lr = F.interpolate(logits_lr, size=masks_hr.shape[-2:], mode="bilinear", align_corners=False)
+                loss_hr = criterion(logits_hr, masks_hr)
+                loss_lr = criterion(logits_lr, masks_lr)
+                loss = loss_hr + lr_supervision_weight * loss_lr
+            running += float(loss.item())
+
+            preds = torch.argmax(logits_hr, dim=1)
+
+            # Ignore index 255 directly on GPU
+            valid = (masks_hr >= 0) & (masks_hr < num_classes)
+            t = masks_hr[valid]
+            p = preds[valid]
+
+            # GPU-vectorized confusion matrix
+            cm = torch.bincount(
+                (t * num_classes + p).view(-1),
+                minlength=num_classes * num_classes
+            ).reshape(num_classes, num_classes)
+            full_confmat += cm
+
+            if is_main and isinstance(iterator, tqdm):
+                iterator.set_postfix(loss=float(loss.item()))
+
+    # Aggregate across all ranks if DDP
+    if is_distributed():
+        dist.all_reduce(full_confmat, op=dist.ReduceOp.SUM)
+
+    avg_loss = running / max(1, len(dl))
+
+    # Convert to CPU once at the end
+    confmat_np = full_confmat.cpu().numpy()
+    metrics = compute_metrics_from_confusion(confmat_np)
+    return avg_loss, metrics
+
+
+def main(cfg_path: str = "configs/FlairHub.yaml"):
+    """Main training function."""
+    logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(levelname)s - %(message)s")
+    cfg = load_config(cfg_path)
+
+    set_seed(cfg.seed)
+
+    # DDP setup
+    setup_distributed()
+    local_rank = get_local_rank()
+    device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
+
+    if cfg.print_device and is_main_process():
+        print(f"Distributed: {is_distributed()} | WORLD_SIZE={os.environ.get('WORLD_SIZE', '1')} | RANK={get_rank()} | LOCAL_RANK={local_rank}")
+        print(f"torch.cuda.is_available(): {torch.cuda.is_available()}")
+        print(f"Device used (this process): {device}")
+
+    # Data
+    dl_train, dl_val, dl_test, train_sampler, val_sampler, test_sampler = make_dataloaders(cfg)
+
+    # Model
+    model = CASWiT(
+        num_head_xa=cfg.cross_attention_heads,
+        num_classes=cfg.num_classes,
+        model_name=cfg.model_name,
+        mlp_ratio=cfg.fusion_mlp_ratio,
+        drop_path=cfg.fusion_drop_path
+    ).to(device)
+
+    # Load pretrained weights
+    if cfg.pretrained_path and cfg.pretrained_path.strip():
+        if is_main_process():
+            print(f"Attempting to load pretrained weights from: {cfg.pretrained_path}")
+        maybe_load_pretrained(model, Path(cfg.pretrained_path))
+    else:
+        if is_main_process():
+            print("No pretrained weights specified. Starting training from scratch.")
+
+    # Wrap with DDP if distributed
+    if is_distributed():
+        model = torch.nn.parallel.DistributedDataParallel(
+            model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=False
+        )
+
+    # Optimizer/Scheduler/Loss
+    criterion = nn.CrossEntropyLoss(ignore_index=cfg.ignore_index)
+    optimizer = optim.AdamW(model.parameters(), lr=cfg.learning_rate)
+    scheduler = CosineAnnealingLR(optimizer, T_max=cfg.num_epochs, eta_min=cfg.eta_min)
+    scaler = GradScaler(enabled=cfg.amp)
+
+    # Logging (wandb only on main)
+    use_wandb = cfg.use_wandb and (wandb is not None) and is_main_process()
+    if use_wandb:
+        setup_wandb_logging(
+            project=cfg.wandb_project,
+            entity=cfg.wandb_entity,
+            run_name=cfg.wandb_run_name,
+            config={
+                "training": {
+                    "num_epochs": cfg.num_epochs,
+                    "learning_rate": cfg.learning_rate,
+                    "batch_size": cfg.batch_size,
+                    "amp": cfg.amp,
+                    "eta_min": cfg.eta_min,
+                },
+                "model": {
+                    "num_classes": cfg.num_classes,
+                    "model_name": cfg.model_name,
+                    "cross_attention_heads": cfg.cross_attention_heads,
+                },
+                "paths": {
+                    "data_path": cfg.data_path,
+                    "save_dir": cfg.save_dir,
+                },
+            },
+            use_wandb=cfg.use_wandb
+        )
+
+    # Train loop with best model tracking
+    Path(cfg.save_dir).mkdir(parents=True, exist_ok=True)
+    best_miou = 0.0
+    best_epoch = 0
+    best_model_path = None
+
+    for epoch in range(cfg.num_epochs):
+        if is_distributed() and train_sampler is not None:
+            train_sampler.set_epoch(epoch)
+
+        if is_main_process():
+            print(f"\nEpoch {epoch + 1}/{cfg.num_epochs}")
+
+        train_loss = train_one_epoch(model, dl_train, device, criterion, optimizer, scaler, cfg.amp, cfg.lr_supervision_weight)
+        val_loss, val_metrics = evaluate(model, dl_val, device, criterion, cfg.num_classes, cfg.lr_supervision_weight, phase_name="Validation")
+
+        current_lr = optimizer.param_groups[0]['lr']
+        current_miou = val_metrics["mIoU"]
+        
+        if is_main_process():
+            log_payload = {
+                "epoch": epoch + 1,
+                "lr": current_lr,
+                "train_loss": train_loss,
+                "val_loss": val_loss,
+                "val_mIoU": current_miou,
+                "val_mF1": val_metrics["mF1"],
+            }
+            print(f"LR={current_lr:.6f} | train_loss={train_loss:.4f} | val_loss={val_loss:.4f} | mIoU={current_miou:.4f}")
+            if use_wandb:
+                log_metrics(log_payload)
+
+            # Save checkpoint for this epoch
+            ckpt_name = f"fusion_hrlr_{cfg.model_name.split('/')[-1]}_lrsupervised_{cfg.batch_size}_epoch_{epoch+1}_head{cfg.cross_attention_heads}.pth"
+            ckpt_path = Path(cfg.save_dir) / ckpt_name
+            torch.save(
+                model.module.state_dict() if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model.state_dict(),
+                str(ckpt_path)
+            )
+
+            # Track best model based on validation mIoU
+            if current_miou > best_miou:
+                best_miou = current_miou
+                best_epoch = epoch + 1
+                best_model_path = ckpt_path
+                # Save best model with special name
+                best_ckpt_name = f"best_model_epoch_{epoch+1}_miou_{current_miou:.4f}.pth"
+                torch.save(
+                    model.module.state_dict() if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model.state_dict(),
+                    str(Path(cfg.save_dir) / best_ckpt_name)
+                )
+                print(f"*** New best model! mIoU: {best_miou:.4f} ***")
+
+        # Step LR scheduler
+        scheduler.step()
+
+    # Final test evaluation with best model
+    if is_main_process():
+        print(f"\n{'='*80}")
+        print(f"Training completed. Best validation mIoU: {best_miou:.4f} at epoch {best_epoch}")
+        print(f"Loading best model for final test evaluation...")
+        print(f"{'='*80}\n")
+
+    # Load best model for final test
+    if best_model_path and best_model_path.exists():
+        state_dict = torch.load(str(best_model_path), map_location=device)
+        if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+            model.module.load_state_dict(state_dict, strict=True)
+        else:
+            model.load_state_dict(state_dict, strict=True)
+        if is_main_process():
+            print(f"Loaded best model from: {best_model_path}\n")
+    
+    # Run final test evaluation
+    test_loss, test_metrics = evaluate(model, dl_test, device, criterion, cfg.num_classes, cfg.lr_supervision_weight, phase_name="Final Test")
+    if is_main_process():
+        test_payload = {
+            "best_epoch": best_epoch,
+            "best_val_mIoU": best_miou,
+            "FINAL_TEST_loss": test_loss,
+            "FINAL_TEST_mIoU": test_metrics["mIoU"],
+            "FINAL_TEST_mF1": test_metrics["mF1"],
+        }
+        print(f"\n{'='*80}")
+        print(f"FINAL TEST RESULTS (with best model from epoch {best_epoch}):")
+        print(f"  Evaluation")
+        print(f"  Test Loss: {test_loss:.4f}")
+        print(f"  Test mIoU: {test_metrics['mIoU']:.4f}")
+        print(f"  Test mF1:  {test_metrics['mF1']:.4f}")
+        print(f"{'='*80}\n")
+        if use_wandb:
+            log_metrics(test_payload)
+    cleanup_distributed()
+
+
+if __name__ == "__main__":
+    import sys
+    cfg_path = sys.argv[1] if len(sys.argv) > 1 else "configs/FlairHub.yaml"
+    main(cfg_path)
+

utils/__init__.py

@@ -0,0 +1,16 @@
+"""
+Utilities for CASWiT training and evaluation.
+"""
+
+from utils.metrics import compute_confusion, compute_metrics_from_confusion
+from utils.logging import setup_wandb_logging, log_metrics
+from utils.attention_viz import viz_cross_attention
+
+__all__ = [
+    'compute_confusion',
+    'compute_metrics_from_confusion',
+    'setup_wandb_logging',
+    'log_metrics',
+    'viz_cross_attention',
+]
+

utils/attention_viz.py

@@ -0,0 +1,225 @@
+"""
+Cross-attention visualization for CASWiT model.
+
+This module provides utilities to visualize cross-attention maps between
+HR and LR branches at different encoder stages.
+"""
+
+from typing import Dict, List, Tuple
+import math
+import torch
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def _to_numpy(img: torch.Tensor) -> np.ndarray:
+    """
+    Convert normalized tensor to numpy array for visualization.
+    
+    Args:
+        img: (1,3,H,W) tensor in normalized space [-1,1]
+        
+    Returns:
+        uint8 HxWx3 array for plotting
+    """
+    x = img.detach().float().cpu()[0]
+    # Undo Normalize(mean=0.5, std=0.5) -> x*0.5 + 0.5
+    x = x * 0.5 + 0.5
+    x = torch.clamp(x, 0, 1)
+    x = (x.permute(1, 2, 0).numpy() * 255.0).astype(np.uint8)
+    return x
+
+
+def _pixel_to_token(x: int, y: int, W_img: int, H_img: int, 
+                    W_tokens: int, H_tokens: int) -> int:
+    """
+    Map HR pixel (x,y) to linear token index for a grid (H_tokens, W_tokens).
+    
+    Uses a ratio-based mapping which is correct for uniform patch embeddings.
+    """
+    tx = min(max(int(math.floor(x * W_tokens / max(W_img, 1))), 0), W_tokens - 1)
+    ty = min(max(int(math.floor(y * H_tokens / max(H_img, 1))), 0), H_tokens - 1)
+    return ty * W_tokens + tx
+
+
+class _CrossAttnTap:
+    """Hook storage for capturing cross-attention weights."""
+    def __init__(self):
+        # Per stage we store: attn (B, N_q, N_k) averaged over heads, and grids
+        self.attn_by_stage: Dict[int, torch.Tensor] = {}
+        self.hr_grid_by_stage: Dict[int, Tuple[int, int]] = {}
+        self.lr_grid_by_stage: Dict[int, Tuple[int, int]] = {}
+        self._handles: List[torch.utils.hooks.RemovableHandle] = []
+
+    def register(self, model) -> None:
+        """Register hooks on model to capture attention weights."""
+        # Locate the list of CrossFusionBlock modules
+        blocks = getattr(model, 'cross_attn_blocks', None)
+        if blocks is None:
+            raise RuntimeError('Model has no attribute cross_attn_blocks')
+        
+        for s, block in enumerate(blocks):
+            # Hook on the CrossFusionBlock to get H/W of inputs (x_hr, x_lr)
+            def fwd_hook(stage_idx: int):
+                def _f(module, inputs, output):
+                    # inputs: (x_hr, x_lr)
+                    x_hr, x_lr = inputs[0], inputs[1]
+                    _, _, Hh, Wh = x_hr.shape
+                    _, _, Hl, Wl = x_lr.shape
+                    self.hr_grid_by_stage[stage_idx] = (Hh, Wh)
+                    self.lr_grid_by_stage[stage_idx] = (Hl, Wl)
+                return _f
+            self._handles.append(block.register_forward_hook(fwd_hook(s)))
+
+            # Hook on the internal nn.MultiheadAttention to grab attn weights
+            mha = getattr(block, 'attn', None)
+            if mha is None:
+                raise RuntimeError(f'CrossFusionBlock at stage {s} has no attn module')
+
+            def attn_hook(stage_idx: int):
+                def _f(module, inputs, output):
+                    # output is a tuple: (attn_out, attn_weights)
+                    if isinstance(output, tuple) and len(output) == 2:
+                        attn_w = output[1]  # shape: (B, N_q, N_k) averaged over heads
+                        self.attn_by_stage[stage_idx] = attn_w.detach()
+                return _f
+            self._handles.append(mha.register_forward_hook(attn_hook(s)))
+
+    def clear(self):
+        """Clear stored attention weights."""
+        self.attn_by_stage.clear()
+        self.hr_grid_by_stage.clear()
+        self.lr_grid_by_stage.clear()
+
+    def remove(self):
+        """Remove all registered hooks."""
+        for h in self._handles:
+            h.remove()
+        self._handles.clear()
+
+
+@torch.no_grad()
+def viz_cross_attention(
+    model: torch.nn.Module,
+    img_hr: torch.Tensor,     # (1,3,H,W) normalized with mean=std=0.5
+    img_lr: torch.Tensor,     # (1,3,h,w) normalized
+    pixel_xy: Tuple[int, int],
+    save_path: str = 'attn_maps.png',
+    overlay_alpha: float = 0.55,
+    dpi: int = 180,
+    show_titles: bool = True,
+):
+    """
+    Visualize cross-attention maps for a given pixel location.
+    
+    Runs a forward pass and saves a multi-panel PNG: one panel per cross-attn stage.
+    The attention is averaged over heads (default behavior of nn.MultiheadAttention).
+    
+    Args:
+        model: CASWiT model (unwrap DDP if needed)
+        img_hr: HR input image [1, 3, H, W]
+        img_lr: LR input image [1, 3, h, w]
+        pixel_xy: (x, y) pixel coordinates in HR image space
+        save_path: Path to save visualization
+        overlay_alpha: Alpha transparency for attention overlay
+        dpi: DPI for saved figure
+        show_titles: Whether to show stage titles
+    """
+    was_training = model.training
+    model.eval()
+
+    # If user passed a DDP-wrapped model, unwrap
+    if hasattr(model, 'module') and not hasattr(model, 'cross_attn_blocks'):
+        model = model.module
+
+    tap = _CrossAttnTap()
+    tap.register(model)
+    tap.clear()
+
+    try:
+        # Forward to populate hooks
+        device = next(model.parameters()).device
+        img_hr = img_hr.to(device)
+        img_lr = img_lr.to(device)
+        _ = model(img_hr, img_lr)
+
+        H_img, W_img = img_hr.shape[-2:]
+        px, py = pixel_xy
+        px = int(np.clip(px, 0, W_img - 1))
+        py = int(np.clip(py, 0, H_img - 1))
+
+        # Prepare base images for overlays (H,W,3) in [0,255]
+        base_hr = _to_numpy(img_hr)
+        base_lr = _to_numpy(img_lr)
+
+        stages = sorted(tap.attn_by_stage.keys())
+        if len(stages) == 0:
+            raise RuntimeError('No attention captured. Ensure a forward pass reached the cross-attention blocks.')
+
+        n = len(stages)
+
+        # Figure with a dedicated column for colorbar
+        fig = plt.figure(figsize=(4.0*n, 4.2), dpi=dpi)
+        gs = fig.add_gridspec(nrows=1, ncols=n+1, width_ratios=[1]*n + [0.04], wspace=0.05)
+
+        axes = [fig.add_subplot(gs[0, i]) for i in range(n)]
+        cax = fig.add_subplot(gs[0, -1])  # Axis reserved for colorbar
+
+        hm = None
+        for i, s in enumerate(stages):
+            attn = tap.attn_by_stage[s]  # (B, N_q, N_k)
+            (Hh, Wh) = tap.hr_grid_by_stage[s]
+            (Hl, Wl) = tap.lr_grid_by_stage[s]
+
+            # Pick batch 0
+            attn0 = attn[0]  # (N_q, N_k)
+            q_idx = _pixel_to_token(px, py, W_img, H_img, Wh, Hh)  # note W first in tokens
+            row = attn0[q_idx]  # (N_k,)
+
+            attn_map = row.view(Hl, Wl)  # reshape to LR grid because K comes from LR branch
+
+            # Normalize for visualization
+            attn_map = attn_map - attn_map.min()
+            denom = float(attn_map.max().item()) if float(attn_map.max().item()) > 0 else 1.0
+            attn_map = attn_map / denom
+
+            # Upsample to LR background size
+            attn_up = F.interpolate(
+                attn_map[None, None, ...], 
+                size=base_lr.shape[:2], 
+                mode='bilinear', 
+                align_corners=False
+            )[0, 0]
+            attn_np = attn_up.detach().cpu().numpy()
+
+            ax = axes[i]
+            ax.imshow(base_lr)
+            hm = ax.imshow(attn_np, cmap='jet', alpha=overlay_alpha, vmin=0.0, vmax=1.0)
+
+            # Approx HR→LR pixel mapping for the marker (simple ratio)
+            hx, hy = base_hr.shape[1], base_hr.shape[0]
+            lx, ly = base_lr.shape[1], base_lr.shape[0]
+            px_lr = int(round(px * lx / max(hx, 1)))
+            py_lr = int(round(py * ly / max(hy, 1)))
+            ax.scatter([px_lr], [py_lr], s=18, c='white', marker='o', 
+                      linewidths=0.5, edgecolors='black')
+
+            if show_titles:
+                ax.set_title(f'Stage {s+1}: HR→LR attn', fontsize=10)
+            ax.set_axis_off()
+
+        # Colorbar in dedicated axis
+        cbar = fig.colorbar(hm, cax=cax)
+        cbar.set_label('Attention')
+
+        # Save PNG
+        fig.savefig(save_path, bbox_inches='tight', format='png')
+        plt.close(fig)
+
+    finally:
+        # Cleanup hooks, restore training mode if needed
+        tap.remove()
+        if was_training:
+            model.train()
+

utils/logging.py

@@ -0,0 +1,40 @@
+"""
+Logging utilities for CASWiT training.
+
+Provides WandB integration and logging helpers.
+"""
+
+from typing import Optional
+try:
+    import wandb
+except ImportError:
+    wandb = None
+
+
+def setup_wandb_logging(project: str, entity: str, run_name: str, config: dict, 
+                       use_wandb: bool = True) -> bool:
+    """
+    Setup Weights & Biases logging.
+    
+    Args:
+        project: WandB project name
+        entity: WandB entity/username
+        run_name: Run name
+        config: Configuration dictionary
+        use_wandb: Whether to use WandB
+        
+    Returns:
+        True if WandB is initialized, False otherwise
+    """
+    if not use_wandb or wandb is None:
+        return False
+    
+    wandb.init(project=project, entity=entity, config=config, name=run_name)
+    return True
+
+
+def log_metrics(metrics: dict, step: Optional[int] = None):
+    """Log metrics to WandB."""
+    if wandb is not None:
+        wandb.log(metrics, step=step)
+

utils/metrics.py

@@ -0,0 +1,47 @@
+"""
+Metrics computation for semantic segmentation.
+
+Provides functions for computing IoU, mIoU, F1 score, and confusion matrices.
+"""
+
+from typing import Dict
+import numpy as np
+from sklearn.metrics import confusion_matrix
+
+
+def compute_confusion(preds: np.ndarray, targets: np.ndarray, num_classes: int) -> np.ndarray:
+    """
+    Compute confusion matrix.
+    
+    Args:
+        preds: Predicted labels
+        targets: Ground truth labels
+        num_classes: Number of classes
+        
+    Returns:
+        Confusion matrix [num_classes, num_classes]
+    """
+    return confusion_matrix(targets.flatten(), preds.flatten(), labels=np.arange(num_classes))
+
+
+def compute_metrics_from_confusion(confmat: np.ndarray) -> Dict[str, np.ndarray]:
+    """
+    Compute IoU, mIoU, and F1 scores from confusion matrix.
+    
+    Args:
+        confmat: Confusion matrix [num_classes, num_classes]
+        
+    Returns:
+        Dictionary with 'mIoU', 'mF1', and 'IoUs' keys
+    """
+    with np.errstate(divide='ignore', invalid='ignore'):
+        intersection = np.diag(confmat)
+        ground_truth_set = confmat.sum(axis=1)
+        predicted_set = confmat.sum(axis=0)
+        union = ground_truth_set + predicted_set - intersection
+        ious = intersection / np.maximum(union, 1)
+        f1s = (2 * intersection) / np.maximum(ground_truth_set + predicted_set, 1)
+        miou = np.nanmean(ious)
+        mf1 = np.nanmean(f1s)
+    return {"mIoU": float(miou), "mF1": float(mf1), "IoUs": ious}
+

weights/CASWiT-Base-SSL_FLAIRHUB_15classes.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a5c54b18a75ca3a462d31a9c70a7ed9575ceacc15e0576a691a3a1db59244f6
+size 1036652553

weights/CASWiT-Base-SSL_URUR_8classes.pth

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

weights/CASWiT-Base_FLAIRHUB_15classes.pth

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

weights/CASWiT-Base_URUR_8classes.pth

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

weights/Swin-Base_FLAIRHUB_15classes.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80aeef0022881619f221b9bf89d9a3174f27ba7f45f701963694e7077a8e4313
+size 489532214