train_seaweed_segmentation.py

"""Train DINOv3 + DeepLabV3+ for seaweed segmentation.

The script supports single GPU, torchrun DDP, AMP, gradient accumulation,
checkpoint resume, and rank-0-only logging/checkpointing.
"""

from __future__ import annotations

import argparse
import json
import os
import random
from datetime import datetime
from pathlib import Path
from typing import Any

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import GradScaler, autocast
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm

from dinov3_deeplabv3plus import DinoV3DeepLabV3Plus, SeaweedSegmentationLoss
from seaweed_segmentation_dataset import SeaweedSegmentationDataset, get_train_transforms, get_val_transforms


def load_json(path: str | Path) -> dict[str, Any]:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def is_dist() -> bool:
    return dist.is_available() and dist.is_initialized()


def rank() -> int:
    return dist.get_rank() if is_dist() else 0


def world_size() -> int:
    return dist.get_world_size() if is_dist() else 1


def is_main_process() -> bool:
    return rank() == 0


def setup_distributed() -> tuple[torch.device, int]:
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1

    if distributed:
        backend = "nccl" if os.name != "nt" else "gloo"
        rank_id = int(os.environ["RANK"])
        size = int(os.environ["WORLD_SIZE"])
        if os.environ.get("USE_LIBUV", "1") == "0":
            from datetime import timedelta

            store = dist.TCPStore(
                os.environ.get("MASTER_ADDR", "127.0.0.1"),
                int(os.environ.get("MASTER_PORT", "29500")),
                world_size=size,
                is_master=rank_id == 0,
                timeout=timedelta(seconds=180),
                use_libuv=False,
            )
            dist.init_process_group(backend=backend, store=store, rank=rank_id, world_size=size)
        else:
            dist.init_process_group(backend=backend)

    if torch.cuda.is_available():
        if distributed:
            torch.cuda.set_device(local_rank)
        device = torch.device(f"cuda:{local_rank}")
    else:
        device = torch.device("cpu")
    return device, local_rank


def cleanup_distributed() -> None:
    if is_dist():
        dist.barrier()
        dist.destroy_process_group()


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def resolve_path(path: str | None, roots: list[str]) -> str | None:
    if not path:
        return path
    candidate = Path(path)
    if candidate.exists():
        return str(candidate)
    for root in roots:
        rooted = Path(root) / path
        if rooted.exists():
            return str(rooted)
    return path


def collate_skip_none(batch: list[dict[str, Any] | None]) -> dict[str, Any]:
    valid = [item for item in batch if item is not None]
    if not valid:
        raise RuntimeError("All samples in this batch failed to load.")
    images = torch.stack([item["image"] for item in valid], dim=0)
    masks = torch.stack([item["mask"] for item in valid], dim=0)
    filenames = [item["filename"] for item in valid]
    return {"image": images, "mask": masks, "filename": filenames}


class SeaweedSegmentationTrainer:
    def __init__(self, config: dict[str, Any], device: torch.device, local_rank: int):
        self.config = config
        self.device = device
        self.local_rank = local_rank
        self.use_amp = bool(config.get("amp", True)) and device.type == "cuda"
        self.grad_accum_steps = int(config.get("grad_accum_steps", 1))

        weight_roots = config.get("weight_search_roots", [])
        config["backbone_weights"] = resolve_path(config.get("backbone_weights"), weight_roots)

        seed_everything(int(config.get("seed", 42)) + rank())

        model = DinoV3DeepLabV3Plus(
            num_classes=config["num_classes"],
            backbone_name=config["backbone_name"],
            pretrained=config["pretrained"],
            weights=config["backbone_weights"],
            use_4channel=config["use_4channel"],
            freeze_backbone=config.get("freeze_backbone", False),
        ).to(device)

        if is_dist():
            self.model = DDP(
                model,
                device_ids=[local_rank] if device.type == "cuda" else None,
                find_unused_parameters=True,
            )
        else:
            self.model = model

        self.criterion = SeaweedSegmentationLoss(
            num_classes=config["num_classes"],
            focal_alpha=config.get("focal_alpha", [1.0, 2.0]),
            focal_gamma=config.get("focal_gamma", 2),
            dice_weight=config.get("dice_weight", 0.5),
            focal_weight=config.get("focal_weight", 1.0),
            background_weight=config.get("background_weight", 1.0),
            foreground_weight=config.get("foreground_weight", 2.0),
        ).to(device)

        self.optimizer = self._create_optimizer()
        self.scheduler = self._create_scheduler()
        self.scaler = GradScaler(enabled=self.use_amp)
        self.train_loader, self.val_loader = self._create_data_loaders()

        self.train_history = {
            "train_loss": [],
            "train_focal": [],
            "train_dice": [],
            "val_loss": [],
            "val_focal": [],
            "val_dice": [],
            "val_iou": [],
            "val_accuracy": [],
        }
        self.best_val_loss = float("inf")
        self.best_val_iou = 0.0
        self.start_epoch = 0

        self.output_dir = Path(config["output_dir"])
        if is_main_process():
            self.output_dir.mkdir(parents=True, exist_ok=True)
            with open(self.output_dir / "config.json", "w", encoding="utf-8") as f:
                json.dump(config, f, indent=2, ensure_ascii=False)

        if config.get("resume"):
            self.load_checkpoint(config["resume"])

    @property
    def raw_model(self) -> nn.Module:
        return self.model.module if isinstance(self.model, DDP) else self.model

    def _create_optimizer(self):
        raw_model = self.raw_model
        if self.config.get("freeze_backbone", False):
            params = raw_model.get_trainable_parameters()
            lr = self.config.get("decoder_lr", 1e-4)
            if is_main_process():
                print(f"Frozen backbone mode. Trainable tensors: {len(params)}")
            return optim.AdamW(params, lr=lr, weight_decay=self.config.get("weight_decay", 1e-4))

        backbone_params = list(raw_model.get_backbone_params())
        decoder_params = list(raw_model.get_decoder_params())
        return optim.AdamW(
            [
                {"params": backbone_params, "lr": self.config.get("backbone_lr", 1e-5)},
                {"params": decoder_params, "lr": self.config.get("decoder_lr", 1e-4)},
            ],
            weight_decay=self.config.get("weight_decay", 1e-4),
        )

    def _create_scheduler(self):
        scheduler_type = self.config.get("scheduler", "cosine")
        if scheduler_type == "cosine":
            return optim.lr_scheduler.CosineAnnealingWarmRestarts(self.optimizer, T_0=10, T_mult=2, eta_min=1e-7)
        if scheduler_type == "step":
            return optim.lr_scheduler.StepLR(self.optimizer, step_size=30, gamma=0.1)
        return None

    def _create_data_loaders(self):
        use_aug = self.config.get("use_data_augmentation", True)
        train_transform = (
            get_train_transforms(self.config["image_size"], self.config["use_4channel"])
            if use_aug
            else get_val_transforms(self.config["image_size"], self.config["use_4channel"])
        )

        train_dataset = SeaweedSegmentationDataset(
            image_dir=self.config["train_image_dir"],
            mask_dir=self.config["train_mask_dir"],
            transform=train_transform,
            target_size=self.config["image_size"],
            use_4channel=self.config["use_4channel"],
        )
        val_dataset = SeaweedSegmentationDataset(
            image_dir=self.config["val_image_dir"],
            mask_dir=self.config["val_mask_dir"],
            transform=get_val_transforms(self.config["image_size"], self.config["use_4channel"]),
            target_size=self.config["image_size"],
            use_4channel=self.config["use_4channel"],
        )

        if len(train_dataset) == 0 or len(val_dataset) == 0:
            raise RuntimeError(
                f"Empty dataset: train={len(train_dataset)}, val={len(val_dataset)}. "
                "Check image/mask directories in the config."
            )

        train_sampler = DistributedSampler(train_dataset, shuffle=True) if is_dist() else None
        val_sampler = DistributedSampler(val_dataset, shuffle=False) if is_dist() else None
        num_workers = int(self.config.get("num_workers", 4))
        loader_kwargs = {
            "batch_size": int(self.config["batch_size"]),
            "num_workers": num_workers,
            "pin_memory": self.device.type == "cuda",
            "collate_fn": collate_skip_none,
            "persistent_workers": num_workers > 0,
        }
        train_loader = DataLoader(
            train_dataset,
            shuffle=train_sampler is None,
            sampler=train_sampler,
            drop_last=True,
            **loader_kwargs,
        )
        val_loader = DataLoader(
            val_dataset,
            shuffle=False,
            sampler=val_sampler,
            drop_last=False,
            **loader_kwargs,
        )
        return train_loader, val_loader

    def calculate_metrics(self, pred, target):
        if isinstance(pred, dict):
            pred = pred["out"]
        pred_classes = torch.argmax(pred, dim=1)
        foreground_pred = pred_classes == 1
        foreground_target = target == 1
        intersection = (foreground_pred & foreground_target).sum().float()
        union = (foreground_pred | foreground_target).sum().float()
        iou = intersection / union.clamp_min(1)
        accuracy = (pred_classes == target).float().mean()
        return {"iou": iou.detach(), "accuracy": accuracy.detach()}

    def _reduce_scalar(self, value: torch.Tensor) -> float:
        value = value.detach().float()
        if is_dist():
            dist.all_reduce(value, op=dist.ReduceOp.SUM)
            value /= world_size()
        return value.item()

    def run_epoch(self, epoch: int, train: bool):
        self.model.train(train)
        loader = self.train_loader if train else self.val_loader
        if train and isinstance(loader.sampler, DistributedSampler):
            loader.sampler.set_epoch(epoch)

        totals = {"total_loss": 0.0, "focal_loss": 0.0, "dice_loss": 0.0, "iou": 0.0, "accuracy": 0.0}
        steps = 0
        desc = f"Epoch {epoch + 1}/{self.config['num_epochs']} - {'Train' if train else 'Val'}"
        iterator = tqdm(loader, desc=desc, disable=not is_main_process())

        if train:
            self.optimizer.zero_grad(set_to_none=True)

        for step, batch in enumerate(iterator):
            max_batches = self.config.get("max_train_batches" if train else "max_val_batches")
            if max_batches is not None and step >= int(max_batches):
                break

            images = batch["image"].to(self.device, non_blocking=True)
            masks = batch["mask"].to(self.device, non_blocking=True)

            with torch.set_grad_enabled(train):
                with autocast(enabled=self.use_amp):
                    outputs = self.model(images)
                    losses = self.criterion(outputs, masks)
                    loss = losses["total_loss"] / self.grad_accum_steps

                if train:
                    self.scaler.scale(loss).backward()
                    should_step = (step + 1) % self.grad_accum_steps == 0 or (step + 1) == len(loader)
                    if should_step:
                        if self.config.get("grad_clip"):
                            self.scaler.unscale_(self.optimizer)
                            torch.nn.utils.clip_grad_norm_(self.raw_model.parameters(), self.config["grad_clip"])
                        self.scaler.step(self.optimizer)
                        self.scaler.update()
                        self.optimizer.zero_grad(set_to_none=True)

            metrics = self.calculate_metrics(outputs, masks)
            totals["total_loss"] += self._reduce_scalar(losses["total_loss"])
            totals["focal_loss"] += self._reduce_scalar(losses["focal_loss"])
            totals["dice_loss"] += self._reduce_scalar(losses["dice_loss"])
            totals["iou"] += self._reduce_scalar(metrics["iou"])
            totals["accuracy"] += self._reduce_scalar(metrics["accuracy"])
            steps += 1

            if is_main_process():
                iterator.set_postfix(
                    loss=f"{totals['total_loss'] / steps:.4f}",
                    iou=f"{totals['iou'] / steps:.4f}",
                    acc=f"{totals['accuracy'] / steps:.4f}",
                )

        return {k: v / max(steps, 1) for k, v in totals.items()}

    def save_checkpoint(self, epoch: int, is_best: bool) -> None:
        if not is_main_process():
            return
        checkpoint = {
            "epoch": epoch,
            "model_state_dict": self.raw_model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "scheduler_state_dict": self.scheduler.state_dict() if self.scheduler else None,
            "scaler_state_dict": self.scaler.state_dict(),
            "train_history": self.train_history,
            "config": self.config,
            "best_val_loss": self.best_val_loss,
            "best_val_iou": self.best_val_iou,
        }
        torch.save(checkpoint, self.output_dir / "latest_checkpoint.pth")
        if is_best:
            torch.save(checkpoint, self.output_dir / "best_checkpoint.pth")
            print(f"Saved best checkpoint: {self.output_dir / 'best_checkpoint.pth'}")

    def load_checkpoint(self, path: str) -> None:
        checkpoint = torch.load(path, map_location=self.device)
        self.raw_model.load_state_dict(checkpoint["model_state_dict"], strict=True)
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        if self.scheduler and checkpoint.get("scheduler_state_dict"):
            self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
        if checkpoint.get("scaler_state_dict"):
            self.scaler.load_state_dict(checkpoint["scaler_state_dict"])
        self.train_history = checkpoint.get("train_history", self.train_history)
        self.best_val_loss = checkpoint.get("best_val_loss", self.best_val_loss)
        self.best_val_iou = checkpoint.get("best_val_iou", self.best_val_iou)
        self.start_epoch = int(checkpoint.get("epoch", -1)) + 1
        if is_main_process():
            print(f"Resumed from {path} at epoch {self.start_epoch}")

    def plot_training_history(self) -> None:
        if not is_main_process() or not self.train_history["train_loss"]:
            return
        epochs = range(1, len(self.train_history["train_loss"]) + 1)
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        axes[0, 0].plot(epochs, self.train_history["train_loss"], label="Train")
        axes[0, 0].plot(epochs, self.train_history["val_loss"], label="Val")
        axes[0, 0].set_title("Loss")
        axes[0, 0].legend()
        axes[0, 1].plot(epochs, self.train_history["train_focal"], label="Train")
        axes[0, 1].plot(epochs, self.train_history["val_focal"], label="Val")
        axes[0, 1].set_title("Focal Loss")
        axes[0, 1].legend()
        axes[1, 0].plot(epochs, self.train_history["train_dice"], label="Train")
        axes[1, 0].plot(epochs, self.train_history["val_dice"], label="Val")
        axes[1, 0].set_title("Dice Loss")
        axes[1, 0].legend()
        axes[1, 1].plot(epochs, self.train_history["val_iou"], label="Val IoU")
        axes[1, 1].plot(epochs, self.train_history["val_accuracy"], label="Val Acc")
        axes[1, 1].set_title("Metrics")
        axes[1, 1].legend()
        for ax in axes.ravel():
            ax.grid(True)
        plt.tight_layout()
        plt.savefig(self.output_dir / "training_history.png", dpi=200, bbox_inches="tight")
        plt.close(fig)

    def train(self):
        if is_main_process():
            effective_batch = self.config["batch_size"] * world_size() * self.grad_accum_steps
            print(f"Device: {self.device}; world_size={world_size()}; AMP={self.use_amp}")
            print(f"Per-GPU batch: {self.config['batch_size']}; effective batch: {effective_batch}")
            print(f"Train samples: {len(self.train_loader.dataset)}; val samples: {len(self.val_loader.dataset)}")
            print(f"Parameters: {sum(p.numel() for p in self.raw_model.parameters()):,}")

        for epoch in range(self.start_epoch, int(self.config["num_epochs"])):
            train_stats = self.run_epoch(epoch, train=True)
            val_stats = self.run_epoch(epoch, train=False)
            if self.scheduler:
                self.scheduler.step()

            self.train_history["train_loss"].append(train_stats["total_loss"])
            self.train_history["train_focal"].append(train_stats["focal_loss"])
            self.train_history["train_dice"].append(train_stats["dice_loss"])
            self.train_history["val_loss"].append(val_stats["total_loss"])
            self.train_history["val_focal"].append(val_stats["focal_loss"])
            self.train_history["val_dice"].append(val_stats["dice_loss"])
            self.train_history["val_iou"].append(val_stats["iou"])
            self.train_history["val_accuracy"].append(val_stats["accuracy"])

            is_best = val_stats["total_loss"] < self.best_val_loss or val_stats["iou"] > self.best_val_iou
            self.best_val_loss = min(self.best_val_loss, val_stats["total_loss"])
            self.best_val_iou = max(self.best_val_iou, val_stats["iou"])

            if is_main_process():
                print(
                    f"Epoch {epoch + 1}: train_loss={train_stats['total_loss']:.4f}, "
                    f"val_loss={val_stats['total_loss']:.4f}, val_iou={val_stats['iou']:.4f}"
                )
            self.save_checkpoint(epoch, is_best)

            if (epoch + 1) % int(self.config.get("plot_interval", 10)) == 0:
                self.plot_training_history()

        self.plot_training_history()
        if is_main_process():
            print(f"Done. best_val_loss={self.best_val_loss:.4f}, best_val_iou={self.best_val_iou:.4f}")


def default_config() -> dict[str, Any]:
    return {
        "train_image_dir": "data/train/images",
        "train_mask_dir": "data/train/masks",
        "val_image_dir": "data/val/images",
        "val_mask_dir": "data/val/masks",
        "num_classes": 2,
        "backbone_name": "dinov3_vitl16",
        "pretrained": True,
        "backbone_weights": "dinov3_vitl16_pretrain_sat493m-eadcf0ff.pth",
        "weight_search_roots": ["D:/DINOv3_pretrained_weights", "D:/DINOv3预训练权重", "dinov3-main"],
        "use_4channel": True,
        "image_size": 256,
        "batch_size": 8,
        "num_epochs": 100,
        "backbone_lr": 1e-5,
        "decoder_lr": 1e-4,
        "weight_decay": 1e-4,
        "scheduler": "cosine",
        "grad_clip": 1.0,
        "grad_accum_steps": 1,
        "amp": True,
        "focal_alpha": [1.0, 3.0],
        "focal_gamma": 2,
        "dice_weight": 0.5,
        "focal_weight": 1.0,
        "background_weight": 1.0,
        "foreground_weight": 3.0,
        "num_workers": 4,
        "output_dir": f"outputs/seaweed_segmentation_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
        "plot_interval": 10,
        "seed": 42,
        "freeze_backbone": False,
        "use_data_augmentation": True,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train seaweed segmentation")
    parser.add_argument("--config", type=str, default=None, help="JSON config path")
    parser.add_argument("--resume", type=str, default=None, help="Checkpoint to resume")
    parser.add_argument("--batch-size", type=int, default=None, help="Override per-GPU batch size")
    parser.add_argument("--epochs", type=int, default=None, help="Override epoch count")
    parser.add_argument("--output-dir", type=str, default=None, help="Override output directory")
    parser.add_argument("--weights", type=str, default=None, help="Override DINOv3 weights path")
    parser.add_argument("--max-train-batches", type=int, default=None, help="Debug limit for train batches per epoch")
    parser.add_argument("--max-val-batches", type=int, default=None, help="Debug limit for val batches per epoch")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    config = default_config()
    if args.config:
        config.update(load_json(args.config))
    if args.resume:
        config["resume"] = args.resume
    if args.batch_size:
        config["batch_size"] = args.batch_size
    if args.epochs:
        config["num_epochs"] = args.epochs
    if args.output_dir:
        config["output_dir"] = args.output_dir
    if args.weights:
        config["backbone_weights"] = args.weights
    if args.max_train_batches is not None:
        config["max_train_batches"] = args.max_train_batches
    if args.max_val_batches is not None:
        config["max_val_batches"] = args.max_val_batches

    device, local_rank = setup_distributed()
    try:
        trainer = SeaweedSegmentationTrainer(config, device, local_rank)
        trainer.train()
    finally:
        cleanup_distributed()


if __name__ == "__main__":
    main()