XNet/train.py

import argparse
import os
import time
from datetime import datetime
from pathlib import Path

import monai
import monai.utils
import swanlab
import torch
from lib.model.model_v4_minute import Wavelet_FFT_SwinUNETR
from monai.metrics import DiceMetric, MeanIoU, HausdorffDistanceMetric
from monai.transforms import (
    Compose, LoadImaged, ScaleIntensityd, RandFlipd, RandRotated, RandRotate90d,
    EnsureChannelFirstd, ToTensord, Resized, Lambdad, RandZoomd, RandShiftIntensityd, RandGaussianNoised,
    RandGaussianSmoothd, RandAdjustContrastd, RandHistogramShiftd,
    RandAxisFlipd, RandCoarseDropoutd,
)
from torch.optim import AdamW
from torch.optim.lr_scheduler import ReduceLROnPlateau

from datasets.PolypDetectionDataset.PolypDetectionDataset import PolypDetectionDataset
from lib.tools.combined_loss import CombinedDiceCEIoULoss


def parse_args():
    parser = argparse.ArgumentParser(description="Polyp Segmentation Model Training Script")

    # ==================== Dataset-related Parameters ====================
    parser.add_argument(
        "--dataset_name",
        type=str,
        required=True,
        help="Dataset name"
    )
    parser.add_argument(
        "--data_root",
        type=str,
        default=r"./data/Polyp-Detection-Dataset",
        help="Root directory path of the dataset"
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=0,
        help="Number of worker processes for data loader"
    )
    parser.add_argument(
        "--pin_memory",
        type=bool,
        default=True,
        help="Whether to enable pinned memory"
    )
    parser.add_argument(
        "--target_spatial_size",
        type=tuple,
        default=(512, 512),
        help="Target spatial size"
    )
    parser.add_argument(
        "--dataset_enhanced",
        type=bool,
        default=True,
        help="Whether to use enhanced data augmentation strategy"
    )

    # ==================== Model-related Parameters ====================
    parser.add_argument(
        "--in_channels",
        type=int,
        default=3,
        help="Number of input image channels"
    )
    parser.add_argument(
        "--out_channels",
        type=int,
        default=1,
        help="Number of output foreground channels"
    )
    parser.add_argument(
        "--feature_size",
        type=int,
        default=48,
        help="Network feature dimension"
    )
    parser.add_argument(
        "--spatial_dims",
        type=int,
        default=2,
        choices=[2, 3],
        help="Spatial dimension (2D or 3D)"
    )
    parser.add_argument(
        "--no_wavelet",
        action="store_false",
        dest="use_wavelet",
        help="Whether to enable wavelet enhancement module"
    )
    parser.add_argument(
        "--wavelet_J",
        type=int,
        default=2,
        help="Wavelet decomposition levels"
    )
    parser.add_argument(
        "--wavelet_wave",
        type=str,
        default="db4",
        help="Wavelet basis type"
    )
    parser.add_argument(
        "--wavelet_reduction",
        type=int,
        default=16,
        help="Wavelet attention compression ratio"
    )
    parser.add_argument(
        "--no_fft",
        action="store_false",
        dest="use_fft",
        help="Whether to enable FFT enhancement module"
    )
    parser.add_argument(
        "--use_v2",
        type=bool,
        default=True,
        help="Whether to enable Swin-UNETR v2 module"
    )

    # ==================== Training-related Parameters ====================
    parser.add_argument(
        "--max_epochs",
        type=int,
        default=1000,
        help="Maximum number of training epochs"
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=4,
        help="Batch size"
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=1e-4,
        help="Learning rate"
    )
    parser.add_argument(
        "--weight_decay",
        type=float,
        default=1e-4,
        help="Weight decay coefficient"
    )

    # ==================== Loss Function Parameters ====================

    parser.add_argument(
        "--dice_weight",
        type=float,
        default=1.0,
        help="Dice loss weight"
    )
    parser.add_argument(
        "--ce_weight",
        type=float,
        default=1.0,
        help="Cross Entropy loss weight"
    )
    parser.add_argument(
        "--iou_weight",
        type=float,
        default=1.0,
        help="IoU loss weight"
    )

    # ==================== SwanLab Parameters ====================
    parser.add_argument(
        "--swanlab_project",
        type=str,
        default="polyp-segmentation-v4_minute",
        help="SwanLab project name"
    )
    parser.add_argument(
        "--swanlab_experiment",
        type=str,
        default=None,
        help="SwanLab experiment name (default uses timestamp)"
    )
    parser.add_argument(
        "--swanlab_log_dir",
        type=str,
        default="./swanlab_log",
        help="SwanLab log directory"
    )

    # ==================== Saving and Loading Parameters ====================
    parser.add_argument(
        "--output_dir",
        type=str,
        default="./outputs_v4_minute",
        help="Directory for saving model checkpoints"
    )
    parser.add_argument(
        "--save_every",
        type=int,
        default=50,
        help="Save model every N epochs"
    )
    # ==================== Early Stopping Parameters ====================
    parser.add_argument(
        "--early_stopping",
        type=bool,
        default=True,
        help="Whether to enable early stopping"
    )
    parser.add_argument(
        "--early_stopping_patience",
        type=int,
        default=100,
        help="Early stopping patience (stop if validation metric doesn't improve for N rounds)"
    )
    parser.add_argument(
        "--early_stopping_min_delta",
        type=float,
        default=1e-4,
        help="Minimum improvement threshold (improvement below this value is considered no improvement)"
    )
    parser.add_argument(
        "--early_stopping_monitor",
        type=str,
        default="dice",
        choices=["dice", "iou", "metric", "loss"],
        help="Metric to monitor for early stopping"
    )
    parser.add_argument(
        "--resume",
        type=str,
        default=None,
        help="Checkpoint path to resume training. If not specified, will automatically load the best Dice model (if exists)"
    )
    parser.add_argument(
        "--no_auto_resume",
        action="store_false",
        dest="auto_resume",
        help="Whether to enable auto-resume functionality (default loads best Dice model)"
    )

    # ==================== Other Parameters ====================
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Training device (cuda or cpu)"
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed"
    )

    return parser.parse_args()


def find_best_checkpoint(args):
    """
    Find the best checkpoint file

    Args:
        args: Command line arguments

    Returns:
        str or None: Best checkpoint path, returns None if not exists
    """
    # Find the best Dice model
    best_dice_path = os.path.join(args.output_dir, f"best_dice_model_{args.dataset_name}.pt")
    if os.path.exists(best_dice_path):
        print(f"Found best Dice model: {best_dice_path}")
        return best_dice_path

    # Find latest checkpoint
    checkpoint_dir = os.path.join(args.output_dir, f"checkpoints_{args.dataset_name}")
    if os.path.exists(checkpoint_dir):
        checkpoints = sorted(
            [f for f in os.listdir(checkpoint_dir) if f.endswith('.pt')],
            key=lambda x: int(x.split('epoch=')[1].split('.')[0]) if 'epoch=' in x else -1,
            reverse=True
        )
        if checkpoints:
            latest_checkpoint = os.path.join(checkpoint_dir, checkpoints[0])
            print(f"Found latest checkpoint: {latest_checkpoint}")
            return latest_checkpoint

    # Find best overall model
    best_metric_path = os.path.join(args.output_dir, f"best_metric_model_{args.dataset_name}.pt")
    if os.path.exists(best_metric_path):
        print(f"Found best overall model: {best_metric_path}")
        return best_metric_path

    return None


def create_enhanced_transforms(target_spatial_size=(512, 512)):
    """
    Enhanced data augmentation strategy

    Includes:
    1. Geometric transformations: flip, rotation, scaling, cropping
    2. Photometric transformations: brightness, contrast, gamma correction
    3. Noise injection: Gaussian noise, low-resolution simulation
    4. Regularization: Coarse Dropout
    """

    def convert_label_to_single_channel(label_tensor):
        """Convert RGB labels to single-channel binary mask"""
        single_channel = label_tensor[0:1, :, :]
        binary_label = (single_channel > 127).float()
        return binary_label

    train_transforms = Compose([
        # ========== Loading and Preprocessing ==========
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        Lambdad(keys=["label"], func=convert_label_to_single_channel),

        # ========== Spatial Transformations ==========
        Resized(keys=["image", "label"], spatial_size=target_spatial_size,
                mode=("bilinear", "nearest")),
        ScaleIntensityd(keys=["image"]),

        # --- Geometric Augmentation ---
        # Random axis flip
        RandAxisFlipd(keys=["image", "label"], prob=0.5),

        # Random rotation (-15° to +15°)
        RandRotated(keys=["image", "label"], range_x=0.15, prob=0.5,
                    keep_size=True, mode=("bilinear", "nearest")),

        # Random 90-degree rotation
        RandRotate90d(keys=["image", "label"], prob=0.5, max_k=2),

        # Random zoom (0.8-1.2x) + cropping
        RandZoomd(keys=["image", "label"], min_zoom=0.8, max_zoom=1.2,
                  prob=0.5, mode=("bilinear", "nearest"), keep_size=True),

        # ========== Photometric Transformations ==========
        # Random brightness adjustment (±20%)
        RandShiftIntensityd(keys=["image"], offsets=(-0.2, 0.2), prob=0.5),

        # Random contrast adjustment (gamma 0.7-1.3)
        RandAdjustContrastd(keys=["image"], gamma=(0.7, 1.3), prob=0.5),

        # Random histogram shift (simulate different staining/lighting conditions)
        RandHistogramShiftd(keys=["image"], num_control_points=(5, 10),
                            prob=0.3),

        # ========== Noise and Quality Degradation ==========
        # Random Gaussian smoothing (simulate blur)
        RandGaussianSmoothd(keys=["image"], sigma_x=(0.5, 1.0),
                            sigma_y=(0.5, 1.0), prob=0.3),

        # Random Gaussian noise
        RandGaussianNoised(keys=["image"], mean=0.0, std=0.05, prob=0.3),

        # Coarse Dropout (occlusion augmentation, improve robustness)
        RandCoarseDropoutd(
            keys=["image"],
            holes=1,
            max_holes=3,
            spatial_size=(32, 32),
            max_spatial_size=(64, 64),
            prob=0.3
        ),

        # ========== Post-processing ==========
        ToTensord(keys=["image", "label"]),
    ])

    return train_transforms


def create_datasets(args):
    """
    Create training and validation datasets

    Args:
        args: Command line arguments

    Returns:
        tuple: (train_dataset, val_dataset)
    """
    print("=" * 60)
    print("正在加载数据集...")
    print("=" * 60)

    def convert_label_to_single_channel(label_tensor):
        """
        Global function: Convert 3-channel RGB labels to 1-channel binary labels (0 or 1)
        Input: label_tensor (shape: [3, H, W], value range 0-255)
        Output: new_tensor (shape: [1, H, W], value range 0 or 1)
        """
        # 1. Extract first channel (R channel)
        single_channel = label_tensor[0:1, :, :]

        # 2. Binarization: pixels greater than 0 are set to 1 (assuming background is pure black 0, polyp is white or other color)
        # This ensures all pixel values can only be 0 or 1, meeting the requirements for out_channels=2
        binary_label = (single_channel > 127).float()

        return binary_label

    # Define training set transformations
    train_transforms = Compose([
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        # Convert labels to single-channel (take first channel or convert to grayscale)
        Lambdad(keys=["label"], func=convert_label_to_single_channel),
        Resized(keys=["image", "label"], spatial_size=args.target_spatial_size, mode=("bilinear", "nearest")),
        ScaleIntensityd(keys=["image"]),
        RandFlipd(keys=["image", "label"], prob=0.5),
        RandRotated(keys=["image", "label"], range_x=0.15, prob=0.5),
    ])
    if args.dataset_enhanced:
        train_transforms = create_enhanced_transforms(args.target_spatial_size)
        print("✓ 使用增强数据增强策略")

    # Define validation set transformations
    val_transforms = Compose([
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        # Convert labels to single-channel (take first channel or convert to grayscale)
        Lambdad(keys=["label"], func=convert_label_to_single_channel),
        Resized(keys=["image", "label"], spatial_size=args.target_spatial_size, mode=("bilinear", "nearest")),
        ScaleIntensityd(keys=["image"]),
    ])

    # Create training dataset
    train_dataset = PolypDetectionDataset(
        root_dir=Path(args.data_root) / args.dataset_name,
        flag='train',
        transform=train_transforms
    )

    # Create validation dataset
    val_dataset = PolypDetectionDataset(
        root_dir=Path(args.data_root) / args.dataset_name,
        flag='val',
        transform=val_transforms
    )

    print(f"✓ Training set size: {len(train_dataset)} samples")
    print(f"✓ Validation set size: {len(val_dataset)} samples")
    print(f"✓ Total samples: {len(train_dataset) + len(val_dataset)} samples")
    print("=" * 60)

    return train_dataset, val_dataset


def create_dataloaders(args, train_dataset, val_dataset):
    """
    Create data loaders

    Args:
        args: Command line arguments
        train_dataset: Training dataset
        val_dataset: Validation dataset

    Returns:
        tuple: (train_loader, val_loader)
    """
    train_loader = monai.data.DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=args.pin_memory,
        drop_last=True
    )

    val_loader = monai.data.DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=args.pin_memory,
        drop_last=False
    )

    print(f"✓ Training data loader: {len(train_loader)} batches")
    print(f"✓ Validation data loader: {len(val_loader)} batches")

    return train_loader, val_loader


def create_model(args):
    """
    Create the model

    Args:
        args: Command line arguments

    Returns:
        torch.nn.Module: Initialized model
    """
    print("\n" + "=" * 60)
    print("Creating model...")

    model = Wavelet_FFT_SwinUNETR(
        in_channels=args.in_channels,
        out_channels=args.out_channels,
        feature_size=args.feature_size,
        spatial_dims=args.spatial_dims,
        wavelet_enhancement=args.use_wavelet,
        wavelet_J=args.wavelet_J,
        wavelet_wave=args.wavelet_wave,
        wavelet_mode='symmetric',
        wavelet_reduction=args.wavelet_reduction,
        fft_enhancement=args.use_fft,
        use_v2=args.use_v2
    )

    # Print model information
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    print(f"\n✓ Total model parameters: {total_params:,}")
    print(f"✓ Trainable parameters: {trainable_params:,}")
    print(f"✓ Using device: {args.device}")
    print("=" * 60)

    return model


def create_loss_function(args):
    """
    Create loss function

    Args:
        args: Command line arguments

    Returns:
        Callable: Loss function
    """
    loss_fn = CombinedDiceCEIoULoss(
        dice_weight=args.dice_weight,
        ce_weight=args.ce_weight,
        iou_weight=args.iou_weight,
        include_background=True,
        to_onehot_y=False,
        softmax=False,
        sigmoid=True,
    )
    return loss_fn


def create_optimizer(args, model):
    """
    Create optimizer

    Args:
        args: Command line arguments
        model: Model

    Returns:
        Optimizer: Optimizer
    """
    optimizer = AdamW(
        model.parameters(),
        lr=args.learning_rate,
        weight_decay=args.weight_decay
    )
    scheduler = ReduceLROnPlateau(
        optimizer,
        mode='min',  # 验证损失越小越好
        factor=0.5,  # 每次乘以 0.5
        patience=20,  # 20 个 epoch 不下降则降低 LR
        threshold=1e-4,  # 最小变化阈值
        cooldown=5,  # 降低 LR 后的冷却期
        min_lr=1e-6  # 学习率下限
    )
    print(f"✓ Optimizer: AdamW")
    print(f"  - Learning rate: {args.learning_rate}")
    print(f"  - Weight decay: {args.weight_decay}")
    print(f"✓ Scheduler: ReduceLROnPlateau")
    print(f"  - Mode: {scheduler.mode}")
    print(f"  - Decay factor: {scheduler.factor}")
    print(f"  - Patience: {scheduler.patience}")
    print(f"  - Minimum change threshold: {scheduler.threshold}")
    print(f"  - Cooldown period: {scheduler.cooldown}")
    print(f"  - Minimum learning rate: {scheduler.min_lrs}")

    return optimizer, scheduler


def setup_swanlab(args):
    """
    Configure SwanLab experiment tracking

    Args:
        args: Command line arguments

    Returns:
        swanlab.Run: SwanLab run object
    """
    # If experiment name is not specified, use timestamp
    if args.swanlab_experiment is None:
        args.swanlab_experiment = "v2_" + args.dataset_name + "_" + datetime.now().strftime("%Y%m%d_%H%M%S")

    # Create log directory
    os.makedirs(args.swanlab_log_dir, exist_ok=True)
    os.makedirs(args.output_dir, exist_ok=True)

    # Initialize SwanLab
    run = swanlab.init(
        project=args.swanlab_project,
        experiment_name=args.swanlab_experiment,
        logdir=args.swanlab_log_dir,
        config=vars(args)
    )

    print(f"\n✓ SwanLab experiment initialized: {args.swanlab_experiment}")
    print(f"  - Project: {args.swanlab_project}")
    print(f"  - Log directory: {args.swanlab_log_dir}")

    return run


def main():
    """
    Main training function
    """
    # ==================== Step 1: Parse arguments ====================
    args = parse_args()

    # Set random seed for reproducibility
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    print("\n" + "=" * 60)
    print("Polyp Segmentation Model Training Started")
    print("=" * 60)
    print(f"Using device: {args.device}")
    print(f"Batch size: {args.batch_size}")
    print(f"Maximum epochs: {args.max_epochs}")
    if args.early_stopping:
        print(
            f"Early stopping: enabled (patience={args.early_stopping_patience}, monitor={args.early_stopping_monitor})")

    # ==================== Step 2: Initialize SwanLab ====================
    run = setup_swanlab(args)

    # ==================== Step 3: Create datasets and data loaders ====================
    train_dataset, val_dataset = create_datasets(args)
    train_loader, val_loader = create_dataloaders(args, train_dataset, val_dataset)

    # ==================== Step 4: Create model, loss function, optimizer ====================
    model = create_model(args)
    model = model.to(args.device)

    loss_function = create_loss_function(args)
    optimizer, scheduler = create_optimizer(args, model)

    # ==================== Step 5: Create evaluation metrics ====================
    dice_metric = DiceMetric(reduction="mean")
    iou_metric = MeanIoU(reduction="mean")
    hd_metric = HausdorffDistanceMetric(reduction="mean")
    # ==================== Step 6: Setup training loop ====================
    best_dice = -1
    best_dice_epoch = -1
    best_metric = -1
    best_metric_epoch = -1
    best_iou = -1
    best_iou_epoch = -1
    epoch_loss_values = []

    dice_metric_values = []
    iou_metric_values = []
    hd_metric_values = []
    start_epoch = 0
    # ==================== Early stopping related variables ====================
    early_stopping_counter = 0
    should_stop = False
    has_restarted = False  # Flag indicating whether it has been restarted once

    # ==================== Step 7: Resume training (if checkpoint exists) ====================
    checkpoint_loaded = False
    checkpoint = None
    if args.resume:
        # User specified checkpoint path
        if not os.path.exists(args.resume):
            raise FileNotFoundError(f"Checkpoint file does not exist: {args.resume}")

        checkpoint_path = args.resume
        print(f"\nResuming training from user-specified checkpoint: {checkpoint_path}")
        checkpoint = torch.load(checkpoint_path, map_location=args.device)
        checkpoint_loaded = True

    elif args.auto_resume:
        # Automatically find best checkpoint
        checkpoint_path = find_best_checkpoint(args)
        if checkpoint_path:
            print(f"\nAuto-resume mode: loading {checkpoint_path}")
            checkpoint = torch.load(checkpoint_path, map_location=args.device)
            checkpoint_loaded = True
        else:
            print("\nNo checkpoints found, starting training from scratch")

    if checkpoint_loaded:
        # Load model weights - supports migration from v1 to v2
        model_dict = model.state_dict()

        # Try to load from checkpoint
        try:
            pretrained_dict = checkpoint["model_state_dict"]
            print("✓ Model weights loaded from training checkpoint")
        except KeyError:
            pretrained_dict = checkpoint
            print("Loading model weights from best Dice or best overall model")

        # Filter and match parameters (handle structural changes from v1->v2)
        matched_params = {}
        unmatched_params = []
        missing_params = []

        for name, param in model_dict.items():
            if name in pretrained_dict:
                pretrained_param = pretrained_dict[name]
                # Check if shape matches
                if param.shape == pretrained_param.shape:
                    matched_params[name] = pretrained_param
                else:
                    unmatched_params.append(f"{name} (shape mismatch: {param.shape} vs {pretrained_param.shape})")
            else:
                missing_params.append(name)

        # Output loading statistics
        print(f"\nWeight loading statistics:")
        print(f"  ✓ Successfully matched parameters: {len(matched_params)}/{len(model_dict)}")
        print(f"  ⚠ Shape mismatched parameters: {len(unmatched_params)}")
        print(f"  ✗ Newly added parameters (randomly initialized): {len(missing_params)}")

        if unmatched_params:
            print(f"\nShape mismatched layers:")
            for info in unmatched_params[:5]:  # Only show first 5
                print(f"  - {info}")
            if len(unmatched_params) > 5:
                print(f"  ... {len(unmatched_params) - 5} more")

        if missing_params:
            print(f"\nNewly added layers (will be randomly initialized):")
            for name in missing_params[:5]:  # Only show first 5
                print(f"  - {name}")
            if len(missing_params) > 5:
                print(f"  ... {len(missing_params) - 5} more")

        # Update pre-trained dictionary
        model_dict.update(matched_params)

        # Load matched parameters
        model.load_state_dict(model_dict, strict=False)
        print(f"\n✓ Model weights loaded (strict mode: False)")
        print("=" * 60)

        if "optimizer_state_dict" in checkpoint:
            # Load optimizer state
            optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
            print("✓ Optimizer state loaded")

        # Load epoch number
        # if "epoch" in checkpoint:
        #     start_epoch = checkpoint["epoch"] + 1  # Start from next epoch
        #     print(f"✓ Training epoch restored to epoch {start_epoch}")

        # Load best metrics
        if "best_dice" in checkpoint:
            best_dice = checkpoint["best_dice"]
            best_dice_epoch = checkpoint["best_dice_epoch"]
            print(f"✓ Best metrics restored: Dice={best_dice:.4f} (Epoch {best_dice_epoch})")

        # Load historical loss and metric values (optional)
        if "epoch_loss_values" in checkpoint:
            epoch_loss_values = checkpoint["epoch_loss_values"]
        if "dice_metric_values" in checkpoint:
            dice_metric_values = checkpoint["dice_metric_values"]

        # Load early stopping state
        if args.early_stopping:
            if "early_stopping_counter" in checkpoint:
                early_stopping_counter = checkpoint["early_stopping_counter"]
                print(f"✓ Early stopping counter restored: {early_stopping_counter}")
            if "should_stop" in checkpoint and checkpoint["should_stop"]:
                should_stop = False  # Even if marked as stopped, allow continued training
                print("✓ Early stopping state reset, can continue training")

        print(f"✓ Training will continue from epoch {start_epoch}")
        print("=" * 60)

    print("\n" + "=" * 60)
    print("Starting training...")
    print("=" * 60)

    start_time = time.time()

    try:
        for epoch in range(start_epoch, run.config.max_epochs):
            # ========== Check early stopping condition ==========
            if should_stop and args.early_stopping:
                print(f"\n{'=' * 60}")
                print(f"Early stopping triggered! Training will terminate early at epoch {epoch + 1}")
                print(f"{'=' * 60}")

                if not has_restarted:
                    # First early stopping: load best weights, restart training
                    print("Early stopping detected, preparing to restart training from best model...")

                    # 1. Find the best Dice model
                    best_checkpoint_path = os.path.join(
                        args.output_dir,
                        f"best_dice_model_{args.dataset_name}.pt"
                    )

                    if os.path.exists(best_checkpoint_path):
                        print(f"Loading best Dice model: {best_checkpoint_path}")
                        checkpoint = torch.load(best_checkpoint_path, map_location=args.device)

                        # 2. Load best weights
                        model.load_state_dict(checkpoint)
                        print("✓ Model weights restored to best state")

                        # 3. Reset optimizer
                        optimizer, scheduler = create_optimizer(args, model)
                        print("✓ Optimizer has been reset")

                        # 4. Reset early stopping counter
                        early_stopping_counter = 0
                        should_stop = False
                        has_restarted = True

                        print("✓ Training restarted from best model")
                        print(f"{'=' * 60}\n")
                        continue  # Skip break, continue to next epoch
                    else:
                        print(f"Warning: Best model file not found {best_checkpoint_path}")
                        print("Will stop training directly")

                # Second early stopping or best model not found: truly stop
                print("Training has been restarted once after early stopping, now stopping training")
                break

            # ========== Training phase ==========
            model.train()
            step = 0
            epoch_loss = 0
            epoch_loss_dice_ce = 0
            epoch_loss_iou = 0

            for batch_data in train_loader:
                step += 1
                inputs = batch_data["image"].to(args.device)
                targets = batch_data["label"].to(args.device)

                optimizer.zero_grad()
                outputs = model(inputs)
                loss, loss_dice_ce, loss_iou = loss_function(outputs, targets)
                loss.backward()
                optimizer.step()

                epoch_loss += loss.item()
                epoch_loss_dice_ce += loss_dice_ce.item()
                epoch_loss_iou += loss_iou.item()

            # If this is the first epoch resumed from checkpoint, print notification
            if epoch == start_epoch and start_epoch > 0:
                print(f"\n✓ Training resumed from epoch {start_epoch}")
            epoch_loss /= step
            epoch_loss_dice_ce /= step
            epoch_loss_iou /= step

            epoch_loss_values.append(epoch_loss)
            print(f"\nEpoch {epoch + 1}/{args.max_epochs} - Training loss: {epoch_loss:.4f}")
            # Log to SwanLab
            swanlab.log({
                "train/loss": epoch_loss,
                "train/loss_dice_ce": epoch_loss_dice_ce,
                "train/loss_iou": epoch_loss_iou,
                "train/lr": optimizer.param_groups[0]['lr'],
            }, step=(epoch + 1))

            # ========== Validation phase ==========
            model.eval()
            val_loss_total = 0
            with torch.no_grad():
                dice_metric.reset()
                iou_metric.reset()
                hd_metric.reset()

                for val_data in val_loader:
                    val_images = val_data["image"].to(args.device)
                    val_labels = val_data["label"].to(args.device)

                    val_outputs = model(val_images)
                    # Calculate validation loss
                    val_loss_batch, _, _ = loss_function(val_outputs, val_labels)
                    val_loss_total += val_loss_batch.item()

                    # Post-processing
                    val_outputs = torch.sigmoid(val_outputs)
                    val_outputs = (val_outputs > 0.5).int()

                    # Calculate Dice score
                    dice_metric(y_pred=val_outputs, y=val_labels)
                    iou_metric(y_pred=val_outputs, y=val_labels)
                    hd_metric(y_pred=val_outputs, y=val_labels)

                # Calculate average validation loss
                val_loss_avg = val_loss_total / len(val_loader)

                # Update learning rate scheduler
                scheduler.step(val_loss_avg)
                current_lr = optimizer.param_groups[0]['lr']

                # Aggregate results
                mean_dice = dice_metric.aggregate().item()
                dice_metric_values.append(mean_dice)
                mean_iou = iou_metric.aggregate().item()
                iou_metric_values.append(mean_iou)
                mean_hd = hd_metric.aggregate().item()
                hd_metric_values.append(mean_hd)

                print(
                    f"Epoch {epoch + 1} - Validation Dice: {mean_dice:.4f}, Validation loss: {val_loss_avg:.4f}, Current LR: {current_lr:.2e}")
                swanlab.log({
                    "val/loss": val_loss_avg,
                    "val/mean_dice": mean_dice,
                    "val/mean_iou": mean_iou,
                    "val/mean_hd": mean_hd,
                    "val/lr": current_lr,
                }, step=(epoch + 1))

                # ========== Early stopping check ==========
                if args.early_stopping:
                    # Get current monitored metric
                    if args.early_stopping_monitor == "dice":
                        current_score = mean_dice
                        best_score = best_dice
                        is_better = current_score > best_score + args.early_stopping_min_delta
                    elif args.early_stopping_monitor == "iou":
                        current_score = mean_iou
                        best_score = best_iou
                        is_better = current_score > best_score + args.early_stopping_min_delta
                    elif args.early_stopping_monitor == "metric":
                        normalized_hd = 1.0 / (1.0 + mean_hd)
                        current_score = 1 * mean_dice + 1 * mean_iou + 1 * normalized_hd
                        best_score = best_metric
                        is_better = current_score > best_score + args.early_stopping_min_delta
                    else:  # loss
                        current_score = -val_loss_avg  # Lower loss is better, so take negative
                        best_score = -min(epoch_loss_values) if epoch_loss_values else float('-inf')
                        is_better = current_score > best_score + args.early_stopping_min_delta

                    # Check if there is improvement
                    if is_better:
                        early_stopping_counter = 0
                        print(
                            f"  ✓ {args.early_stopping_monitor.upper()} metric improved: {current_score:.4f} > {best_score:.4f}")
                    else:
                        early_stopping_counter += 1
                        print(
                            f"  ⚠ {args.early_stopping_monitor.upper()} metric did not improve, counter: {early_stopping_counter}/{args.early_stopping_patience}")

                        # Check if early stopping should be triggered
                        if early_stopping_counter >= args.early_stopping_patience:
                            should_stop = True

                # Save best Dice model
                if mean_dice > best_dice:
                    best_dice = mean_dice
                    best_dice_epoch = epoch + 1
                    checkpoint_path = os.path.join(args.output_dir, f"best_dice_model_{args.dataset_name}.pt")
                    Path(checkpoint_path).parent.mkdir(parents=True, exist_ok=True)
                    torch.save(model.state_dict(), checkpoint_path)
                    print(
                        f"✓ Found better Dice model! Dice: {mean_dice:.4f}, IoU: {mean_iou:.4f}, HD: {mean_hd:.4f}, saved to {checkpoint_path}")
                # Save best IoU model
                if mean_iou > best_iou:
                    best_iou = mean_iou
                    best_iou_epoch = epoch + 1
                    checkpoint_path = os.path.join(args.output_dir, f"best_iou_model_{args.dataset_name}.pt")
                    Path(checkpoint_path).parent.mkdir(parents=True, exist_ok=True)
                    torch.save(model.state_dict(), checkpoint_path)
                    print(
                        f"✓ Found better IoU model! IoU: {mean_iou:.4f}, Dice: {mean_dice:.4f}, HD: {mean_hd:.4f}, saved to {checkpoint_path}"
                    )
                # Save best overall model
                normalized_hd = 1.0 / (1.0 + mean_hd)
                mean_metric = (
                        1 * mean_dice +
                        1 * mean_iou +
                        1 * normalized_hd
                )
                if mean_metric > best_metric:
                    best_metric = mean_metric
                    best_metric_epoch = epoch + 1
                    checkpoint_path = os.path.join(args.output_dir, f"best_metric_model_{args.dataset_name}.pt")
                    Path(checkpoint_path).parent.mkdir(parents=True, exist_ok=True)
                    torch.save(model.state_dict(), checkpoint_path)
                    print(
                        f"✓ Found better overall model! Overall score: {mean_metric:.4f}, Dice: {mean_dice:.4f}, IoU: {mean_iou:.4f}, HD: {mean_hd:.4f}, saved to {checkpoint_path}"
                    )
            # Periodically save checkpoint
            if (epoch + 1) % args.save_every == 0:
                checkpoint_path = os.path.join(args.output_dir, f"checkpoints_{args.dataset_name}",
                                               f"checkpoint_epoch={epoch}.pt")
                Path(checkpoint_path).parent.mkdir(parents=True, exist_ok=True)
                torch.save({
                    "epoch": epoch,
                    "model_state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "best_dice": best_dice,
                    "best_dice_epoch": best_dice_epoch,
                    "epoch_loss_values": epoch_loss_values,
                    "dice_metric_values": dice_metric_values,
                    "iou_metric_values": iou_metric_values,
                    "hd_metric_values": hd_metric_values,
                    "best_metric": best_metric,
                    "best_metric_epoch": best_metric_epoch,
                    "best_iou": best_iou,
                    "best_iou_epoch": best_iou_epoch,
                    "early_stopping_counter": early_stopping_counter,
                    "should_stop": should_stop
                }, checkpoint_path)
                print(f"✓ Checkpoint saved: {checkpoint_path}")

    except KeyboardInterrupt:
        print("\nTraining interrupted by user")
    finally:
        end_time = time.time()
        training_time = end_time - start_time

        print("\n" + "=" * 60)
        print("Training completed!")
        print(f"Total training time: {training_time / 3600:.2f} hours")
        print(f"Best validation Dice: {best_dice:.4f} (Epoch {best_dice_epoch})")
        print("=" * 60)

        # Close SwanLab
        swanlab.finish()
        print("✓ SwanLab experiment saved")


if __name__ == "__main__":
    main()