agrobot_base/Python/OAK/datasets/_8_train_segformer.py

#python _8_train_segformer_b3.py --epochs 180 --batch 2 --lr 3e-5 --wd 0.01 --num_workers 4 --amp --amp_val --grad_accum 2 --class_weights auto --main_class navegavel --resume

# _8_train_segformer_b3.py  (PATCH)
import os

# (Opcional) ajuda com fragmentação em algumas máquinas.
# Idealmente isso deveria vir ANTES de importar torch, mas já ajuda quando setado fora também.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128")

import json
import math
import time
import argparse
from typing import Dict, List, Optional, Tuple, Any

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

# >>> troque AMP deprecated
from torch.amp import autocast, GradScaler

from transformers import SegformerForSemanticSegmentation

# >>>>>>>>> AJUSTE AQUI <<<<<<<<<
from _8_train_fastscnn_v2 import ROISegDataset  # troque se necessário


def set_seed(seed: int = 42):
    import random
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def _ids_by_names(class_names: Any, wanted_names: List[str]) -> List[int]:
    if class_names is None:
        return []
    wanted_names = [str(x).lower() for x in wanted_names]

    if isinstance(class_names, list):
        low = [c.lower() for c in class_names]
        return [low.index(n) for n in wanted_names if n in low]

    if isinstance(class_names, dict):
        if all(isinstance(k, int) for k in class_names.keys()):
            inv = {str(v).lower(): int(k) for k, v in class_names.items()}
            return [inv[n] for n in wanted_names if n in inv]
        inv = {str(k).lower(): int(v) for k, v in class_names.items()}
        return [inv[n] for n in wanted_names if n in inv]

    return []


@torch.no_grad()
def update_confusion_matrix(cm: torch.Tensor, preds: torch.Tensor, labels: torch.Tensor, num_classes: int, ignore_index: int = 255):
    preds = preds.view(-1)
    labels = labels.view(-1)

    valid = labels != ignore_index
    preds = preds[valid]
    labels = labels[valid]

    idx = labels * num_classes + preds
    bins = torch.bincount(idx, minlength=num_classes * num_classes)
    cm += bins.view(num_classes, num_classes)


@torch.no_grad()
def compute_iou_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> Tuple[float, List[float]]:
    cm = cm.float()
    tp = torch.diag(cm)
    fp = cm.sum(0) - tp
    fn = cm.sum(1) - tp
    denom = tp + fp + fn + eps
    iou = (tp / denom).cpu().tolist()
    miou = float(np.mean(iou))
    return miou, iou


@torch.no_grad()
def compute_pixel_acc_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> float:
    cm = cm.float()
    acc = (torch.diag(cm).sum() / (cm.sum() + eps)).item()
    return acc


IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]
NORM_MEAN = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
NORM_STD = torch.tensor(IMAGENET_STD).view(3, 1, 1).clamp_min(1e-6)
def normalize_img(img: torch.Tensor) -> torch.Tensor:
    return (img - NORM_MEAN.to(img.device)) / NORM_STD.to(img.device)


def default_collate(batch):
    imgs = []
    masks = []
    for item in batch:
        if isinstance(item, dict):
            img = item["image"]
            mask = item["mask"]
        else:
            img, mask = item

        if isinstance(img, np.ndarray):
            img = torch.from_numpy(img)
        if isinstance(mask, np.ndarray):
            mask = torch.from_numpy(mask)

        if img.ndim == 3 and img.shape[-1] == 3:
            img = img.permute(2, 0, 1)

        if img.dtype != torch.float32:
            img = img.float()
            if img.max() > 1.5:
                img = img / 255.0

        imgs.append(img)
        masks.append(mask.long())

    return torch.stack(imgs, 0), torch.stack(masks, 0)


def estimate_class_weights(ds: ROISegDataset, num_classes: int, ignore_index: int = 255, max_samples: int = 800) -> torch.Tensor:
    n = min(len(ds), max_samples)
    idxs = np.random.choice(len(ds), size=n, replace=False)

    counts = np.zeros(num_classes, dtype=np.float64)
    for i in idxs:
        item = ds[i]
        mask = item["mask"] if isinstance(item, dict) else item[1]
        m = mask.cpu().numpy() if isinstance(mask, torch.Tensor) else np.array(mask)

        m = m.reshape(-1)
        m = m[m != ignore_index]
        if m.size == 0:
            continue

        counts += np.bincount(m, minlength=num_classes)[:num_classes]

    freq = counts / (counts.sum() + 1e-12)
    freq = np.clip(freq, 1e-12, 1.0)
    weights = 1.0 / np.log(1.02 + freq)
    weights = weights / weights.mean()
    return torch.tensor(weights, dtype=torch.float32)


def save_checkpoint(path: str,
                    model: nn.Module,
                    optimizer: torch.optim.Optimizer,
                    scaler: Optional[GradScaler],
                    epoch: int,
                    best_miou: float,
                    best_main_iou: float,
                    extra: Optional[Dict[str, Any]] = None):
    ckpt = {
        "epoch": epoch,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "best_miou": best_miou,
        "best_main_iou": best_main_iou,
    }
    if scaler is not None:
        ckpt["scaler"] = scaler.state_dict()
    if extra:
        ckpt["extra"] = extra
    torch.save(ckpt, path)


def load_checkpoint(path: str,
                    model: nn.Module,
                    optimizer: Optional[torch.optim.Optimizer] = None,
                    scaler: Optional[GradScaler] = None,
                    map_location: str = "cpu") -> Dict[str, Any]:
    ckpt = torch.load(path, map_location=map_location, weights_only=False)
    model.load_state_dict(ckpt["model"], strict=True)
    if optimizer is not None and "optimizer" in ckpt:
        optimizer.load_state_dict(ckpt["optimizer"])
    if scaler is not None and "scaler" in ckpt:
        scaler.load_state_dict(ckpt["scaler"])
    return ckpt


def run_one_epoch(model: nn.Module,
                  loader: DataLoader,
                  optimizer: Optional[torch.optim.Optimizer],
                  device: torch.device,
                  num_classes: int,
                  ignore_index: int,
                  criterion: nn.Module,
                  amp: bool,
                  scaler: Optional[GradScaler],
                  train: bool,
                  grad_accum: int = 1) -> Dict[str, Any]:

    model.train(train)

    total_loss = 0.0
    cm = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device)

    t0 = time.time()
    n_batches = 0

    # >>>>> ESSENCIAL: desliga grad no val
    with torch.set_grad_enabled(train):
        if train and optimizer is not None:
            optimizer.zero_grad(set_to_none=True)

        for step, (imgs, masks) in enumerate(loader):
            imgs = imgs.to(device, non_blocking=True)
            masks = masks.to(device, non_blocking=True)

            imgs = normalize_img(imgs)

            # AMP tanto em train quanto (opcionalmente) em val
            with autocast(device_type="cuda", enabled=amp and device.type == "cuda"):
                out = model(pixel_values=imgs)
                logits = out.logits

                if logits.shape[-2:] != masks.shape[-2:]:
                    logits = torch.nn.functional.interpolate(
                        logits, size=masks.shape[-2:], mode="bilinear", align_corners=False
                    )

                loss = criterion(logits, masks)
                if train and grad_accum > 1:
                    loss = loss / grad_accum

            if train and optimizer is not None:
                if amp and scaler is not None and device.type == "cuda":
                    scaler.scale(loss).backward()
                    if ((step + 1) % grad_accum) == 0:
                        scaler.step(optimizer)
                        scaler.update()
                        optimizer.zero_grad(set_to_none=True)
                else:
                    loss.backward()
                    if ((step + 1) % grad_accum) == 0:
                        optimizer.step()
                        optimizer.zero_grad(set_to_none=True)

            total_loss += float(loss.item()) * (grad_accum if train and grad_accum > 1 else 1.0)
            n_batches += 1

            preds = torch.argmax(logits, dim=1)
            update_confusion_matrix(cm, preds, masks, num_classes=num_classes, ignore_index=ignore_index)

    dt = time.time() - t0
    avg_loss = total_loss / max(n_batches, 1)
    miou, iou_per_class = compute_iou_from_cm(cm)
    acc = compute_pixel_acc_from_cm(cm)

    return {"loss": avg_loss, "miou": miou, "iou_per_class": iou_per_class, "acc": acc, "time_s": dt}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="config.json")
    parser.add_argument("--epochs", type=int, default=120)
    parser.add_argument("--batch", type=int, default=1)  # <<< default seguro pra 8GB
    parser.add_argument("--lr", type=float, default=6e-5)
    parser.add_argument("--wd", type=float, default=0.01)
    parser.add_argument("--num_workers", type=int, default=4)

    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--resume", action="store_true")

    parser.add_argument("--amp", action="store_true")
    parser.add_argument("--amp_val", action="store_true")  # <<< novo: AMP no val
    parser.add_argument("--grad_accum", type=int, default=1)  # <<< novo

    parser.add_argument("--grad_ckpt", action="store_true")  # <<< novo: checkpointing
    parser.add_argument("--ignore_index", type=int, default=255)
    parser.add_argument("--class_weights", type=str, default="auto")
    parser.add_argument("--main_class", type=str, default=None)
    parser.add_argument("--es_classes", type=str, default="")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--norm_stats", type=str, default=None, help="Caminho para JSON com mean/std por canal (ex: norm_stats.json).")
    args = parser.parse_args()

    set_seed(args.seed)

    with open(args.config, "r") as f:
        config = json.load(f)

    MODELO = config["camera"]
    MODEL_NAME = config["model_name"]
    RESOLUCAO = config["resolucao"]
    ROI_INICIO = config["roi_inicio"]
    ROI_TAMANHO = config["roi_tamanho"]
    MAIN_CLASS_NAME = str(config.get("main_class_name", "erva")).lower()
    BACKBONE = config["backbone"]
    if args.main_class is not None:
        MAIN_CLASS_NAME = args.main_class.lower()

    dataset_path = os.path.join(MODELO, "dataset")
    labelmap_path = os.path.join(dataset_path, "labelmap.txt")

    save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME)
    os.makedirs(save_path, exist_ok=True)
    last_ckpt_path = os.path.join(save_path, "last.pt")
    best_miou_path = os.path.join(save_path, "best_miou.pt")
    best_main_path = os.path.join(save_path, "best_main.pt")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    # Datasets
    ds_train = ROISegDataset(
        os.path.join(dataset_path, "split", "train"),
        save_path, ROI_INICIO, ROI_TAMANHO,
        RESOLUCAO[0], RESOLUCAO[1], labelmap_path
    )
    ds_val = ROISegDataset(
        os.path.join(dataset_path, "split", "val"),
        save_path, ROI_INICIO, ROI_TAMANHO,
        RESOLUCAO[0], RESOLUCAO[1], labelmap_path
    )

    CLASS_NAMES = getattr(ds_train, "classes", None)
    if CLASS_NAMES is None:
        raise RuntimeError("ROISegDataset precisa expor .classes (list ou dict).")

    if isinstance(CLASS_NAMES, list):
        num_classes = len(CLASS_NAMES)
        class_id_by_name = {n.lower(): i for i, n in enumerate(CLASS_NAMES)}
        class_name_by_id = {i: n for i, n in enumerate(CLASS_NAMES)}
    elif isinstance(CLASS_NAMES, dict):
        if all(isinstance(k, int) for k in CLASS_NAMES.keys()):
            num_classes = len(CLASS_NAMES)
            class_name_by_id = {int(k): str(v) for k, v in CLASS_NAMES.items()}
            class_id_by_name = {str(v).lower(): int(k) for k, v in CLASS_NAMES.items()}
        else:
            class_id_by_name = {str(k).lower(): int(v) for k, v in CLASS_NAMES.items()}
            num_classes = len(class_id_by_name)
            class_name_by_id = {v: k for k, v in class_id_by_name.items()}
    else:
        raise RuntimeError("Formato de .classes não reconhecido.")

    main_class_id = class_id_by_name.get(MAIN_CLASS_NAME, None)
    if main_class_id is None:
        print(f"[WARN] main_class_name='{MAIN_CLASS_NAME}' não encontrado. best_main_iou usa mIoU.")
    else:
        print(f"Main class: '{MAIN_CLASS_NAME}' -> id={main_class_id}")

    es_names = [x.strip().lower() for x in args.es_classes.split(",") if x.strip()]
    ES_CLASS_IDS = _ids_by_names(CLASS_NAMES, es_names)

    dl_train = DataLoader(
        ds_train, batch_size=args.batch, shuffle=True,
        num_workers=args.num_workers, pin_memory=True,
        collate_fn=default_collate, drop_last=True
    )
    dl_val = DataLoader(
        ds_val, batch_size=1, shuffle=False,   # <<< val com batch 1 é mais estável
        num_workers=max(2, args.num_workers // 2), pin_memory=True,
        collate_fn=default_collate, drop_last=False
    )

    model = SegformerForSemanticSegmentation.from_pretrained(
        BACKBONE,
        num_labels=num_classes,
        ignore_mismatched_sizes=True,
        use_safetensors=True
    )

    if args.grad_ckpt:
        try:
            # Nem toda versão/classe suporta
            model.gradient_checkpointing_enable()
            print("[OK] gradient checkpointing enabled")
        except Exception as e:
            print(f"[WARN] gradient checkpointing não suportado aqui ({type(e).__name__}: {e}). Seguindo sem.")

    model.to(device)

    # Loss weights
    if args.class_weights.lower() == "none":
        weights = None
    elif args.class_weights.lower() == "auto":
        w = estimate_class_weights(ds_train, num_classes=num_classes, ignore_index=args.ignore_index)
        weights = w.to(device)
        print("Class weights (auto):", w.cpu().numpy().round(3).tolist())
    else:
        parts = [float(x) for x in args.class_weights.split(",")]
        if len(parts) != num_classes:
            raise ValueError(f"class_weights manual precisa ter {num_classes} valores, recebeu {len(parts)}.")
        weights = torch.tensor(parts, dtype=torch.float32, device=device)

    criterion = nn.CrossEntropyLoss(weight=weights, ignore_index=args.ignore_index)

    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wd)

    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=0.5, patience=6, threshold=1e-4, verbose=True
    )

    scaler = GradScaler(enabled=args.amp and device.type == "cuda")



    # ==========================
    # Normalizador (fixo ou dinâmico)
    # ==========================
    global NORM_MEAN, NORM_STD
    norm_stats = None

    # Caminho padrão: dentro do dataset, nome do arquivo de stats
    # (ajusta aqui pro nome que você realmente usou: norm_stats.json, por ex.)
    norm_stats_path = os.path.join(save_path, "norm_stats.json")
    if args.norm_stats is not None:
        norm_stats_path = args.norm_stats

    if norm_stats_path is not None and os.path.exists(norm_stats_path):
        with open(norm_stats_path, "r", encoding="utf-8") as f:
            norm_stats = json.load(f)

        stats_channels = norm_stats.get("channels", [])
        stats_mean = norm_stats.get("mean", [])
        stats_std  = norm_stats.get("std", [])

        print(f"[NORM] usando stats fixos de: {norm_stats_path}")
        print(f"[NORM] channels={stats_channels}")
        print(f"[NORM] mean={stats_mean}")
        print(f"[NORM] std ={stats_std}")

        # Garante que temos pelo menos R,G,B
        idx_by_name = {name: i for i, name in enumerate(stats_channels)}
        required = ["R", "G", "B"]
        if not all(ch in idx_by_name for ch in required):
            print("[NORM] AVISO: norm_stats não contém todos os canais R,G,B. Mantendo normalize imagenet.")
        else:
            NORM_MEAN = torch.tensor(stats_mean, dtype=torch.float32, device=device).view(3, 1, 1)
            NORM_STD = torch.tensor(stats_std, dtype=torch.float32, device=device).view(3, 1, 1).clamp_min(1e-6)

            print("[NORM] Normalização fixa por canal ativada para [R,G,B].")
    else:
        if norm_stats_path:
            print(f"[NORM] Caminho de norm_stats não encontrado: {norm_stats_path}. Usando normalize imagenet.")
        else:
            print("[NORM] norm_stats não informado. Usando normalize imagenet.")



    start_epoch = 1
    best_miou = -1.0
    best_main_iou = -1.0

    if args.resume and os.path.exists(best_miou_path):
        ckpt = load_checkpoint(best_miou_path, model, optimizer, scaler=scaler, map_location="cpu")
        start_epoch = int(ckpt["epoch"]) + 1
        best_miou = float(ckpt.get("best_miou", -1.0))
        best_main_iou = float(ckpt.get("best_main_iou", -1.0))
        print(f"[RESUME] epoch={start_epoch} best_miou={best_miou:.4f} best_main_iou={best_main_iou:.4f}")

    def pretty_iou(iou_list):
        return " | ".join([f"{class_name_by_id.get(cid,cid)}:{v:.3f}" for cid, v in enumerate(iou_list)])

    for epoch in range(start_epoch, args.epochs + 1):
        lr_now = optimizer.param_groups[0]["lr"]
        print(f"\n==== Epoch {epoch}/{args.epochs} | lr={lr_now:.2e} ====")

        if device.type == "cuda":
            torch.cuda.empty_cache()

        tr = run_one_epoch(
            model=model, loader=dl_train, optimizer=optimizer,
            device=device, num_classes=num_classes, ignore_index=args.ignore_index,
            criterion=criterion, amp=args.amp, scaler=scaler, train=True,
            grad_accum=max(1, args.grad_accum)
        )

        if device.type == "cuda":
            torch.cuda.empty_cache()

        va = run_one_epoch(
            model=model, loader=dl_val, optimizer=None,
            device=device, num_classes=num_classes, ignore_index=args.ignore_index,
            criterion=criterion, amp=args.amp_val, scaler=None, train=False,
            grad_accum=1
        )

        scheduler.step(va["loss"])

        main_iou = va["miou"] if main_class_id is None else va["iou_per_class"][main_class_id]

        print(f"TRAIN: loss={tr['loss']:.4f} acc={tr['acc']:.4f} miou={tr['miou']:.4f} (t={tr['time_s']:.1f}s)")
        print(f"VAL  : loss={va['loss']:.4f} acc={va['acc']:.4f} miou={va['miou']:.4f} main_iou={main_iou:.4f} (t={va['time_s']:.1f}s)")
        print("IoU per class:", pretty_iou(va["iou_per_class"]))

        if ES_CLASS_IDS:
            mini = " | ".join([f"{class_name_by_id[cid]}:{va['iou_per_class'][cid]:.3f}" for cid in ES_CLASS_IDS])
            print("ES classes:", mini)

        save_checkpoint(
            last_ckpt_path, model, optimizer, scaler=scaler,
            epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou,
            extra={"val_loss": va["loss"], "val_miou": va["miou"], "val_main_iou": float(main_iou)}
        )

        if args.save_every > 0 and (epoch % args.save_every == 0):
            save_checkpoint(
                os.path.join(save_path, f"epoch_{epoch:04d}.pt"),
                model, optimizer, scaler=scaler,
                epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou
            )

        if va["miou"] > best_miou:
            best_miou = va["miou"]
            save_checkpoint(
                best_miou_path, model, optimizer, scaler=scaler,
                epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou
            )
            print(f"[BEST mIoU] {best_miou:.4f} -> saved: {best_miou_path}")

        if float(main_iou) > best_main_iou:
            best_main_iou = float(main_iou)
            save_checkpoint(
                best_main_path, model, optimizer, scaler=scaler,
                epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou
            )
            print(f"[BEST MAIN] {best_main_iou:.4f} -> saved: {best_main_path}")


if __name__ == "__main__":
    main()