agrobot_base/Python/OAK/datasets/_5_train_fastscnn.py

import json
import os
import time
import argparse
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from fast_scnn import FastSCNN
from roi_seg_dataset import ROISegDataset
import matplotlib.pyplot as plt

# ⚙️ Configurações
with open("config.json", "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 = config["main_class_name"]
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME)
dataset_path = os.path.join(MODELO, "dataset")
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
batch_size = 8
num_workers = 4

# ---- Helpers de métricas ----
@torch.no_grad()
def confmat_update(confmat, pred, target, num_classes, ignore_index=None):
    # pred, target: (B,H,W)
    if ignore_index is not None:
        mask = target != ignore_index
        target = target[mask]
        pred   = pred[mask]
    k = (target * num_classes + pred).to(torch.int64)
    binc = torch.bincount(k, minlength=num_classes**2)
    confmat += binc.reshape(num_classes, num_classes)
    return confmat

def metrics_from_confmat(confmat, main_class_id=None):
    # confmat: CxC
    cm = confmat.float()
    tp = torch.diag(cm)
    fp = cm.sum(0) - tp
    fn = cm.sum(1) - tp
    denom_iou = tp + fp + fn + 1e-7
    iou_per_class = tp / denom_iou
    miou = iou_per_class.mean().item()
    pix_acc = tp.sum() / (cm.sum() + 1e-7)

    main_class_metrics = None
    if main_class_id is not None and 0 <= main_class_id < cm.shape[0]:
        p = tp[main_class_id] / (tp[main_class_id] + fp[main_class_id] + 1e-7)
        r = tp[main_class_id] / (tp[main_class_id] + fn[main_class_id] + 1e-7)
        f1 = 2 * p * r / (p + r + 1e-7)
        main_class_metrics = {
            "precision": p.item(),
            "recall": r.item(),
            "f1": f1.item(),
            "iou": iou_per_class[main_class_id].item(),
        }
    return {
        "miou": miou,
        "pixel_acc": pix_acc.item(),
        "iou_per_class": iou_per_class.cpu().tolist(),
        "main_class": main_class_metrics
    }

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

    # --- Dataset ---
    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
    )

    dl_train = DataLoader(ds_train, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=True)
    dl_val   = DataLoader(ds_val,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)

    # Detecta automaticamente o ID da classe ERVA
    main_class_id = None
    try:
        if hasattr(ds_train, "classes") and isinstance(ds_train.classes, dict):
            for k, v in ds_train.classes.items():
                if isinstance(v, str) and MAIN_CLASS_NAME in v.lower():
                    main_class_id = k
                    break
        elif isinstance(ds_train.classes, (list, tuple)):
            main_class_id = next((i for i, c in enumerate(ds_train.classes) if isinstance(c, str) and MAIN_CLASS_NAME in c.lower()), None)

        if main_class_id is not None:
            print(f"🌿 Classe PRIMARIA detectada: id={main_class_id}, nome='{ds_train.classes[main_class_id]}'")
        else:
            print("⚠️ Classe PRIMARIA não encontrada; métricas específicas da classe primaria serão puladas.")
    except Exception as e:
        print(f"⚠️ Erro ao detectar classe PRIMARIA: {e}")

    num_classes = len(ds_train.classes)

    # --- Modelo / Otimizador / Schedulers ---
    model = FastSCNN(num_classes=num_classes).to(device)
    criterion = nn.CrossEntropyLoss(ignore_index=ds_train.ignore_id)
    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)

    # Scheduler inteligente: começa em Cosine, muda pra Plateau se travar
    min_lr = getattr(args, "min_lr", 1e-6)
    plateau_factor = getattr(args, "plateau_factor", 0.5)
    plateau_patience = getattr(args, "plateau_patience", 6)  # épocas sem melhora antes de trocar
    plateau_cooldown = getattr(args, "plateau_cooldown", 1)

    cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs, eta_min=min_lr)
    plateau = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=plateau_factor,
        patience=plateau_patience, cooldown=plateau_cooldown,
        min_lr=min_lr, verbose=True
    )
    active_sched = "cosine"

    scaler = torch.cuda.amp.GradScaler(enabled=args.amp)

    start_epoch = 1
    best_val_loss = float("inf")
    best_main_class_f1  = -1.0
    train_loss_history, val_loss_history, lr_history = [], [], []
    f1_history, miou_history = [], []

    # --- no topo (config) ---
    patience_loss = 12          # ligeiramente > plateau_patience + 2
    patience_f1   = 6           # deixa o F1 respirar
    delta_f1_min  = 0.0015      # ignora ruído
    grace_after_switch = 4      # épocas de graça após mudar pro Plateau

    no_imp_loss = 0
    no_imp_f1   = 0
    epochs_since_switch = 0
    active_sched = "cosine"  # como já está

    # --- Checkpoint ---
    if args.checkpoint and os.path.exists(args.checkpoint):
        print(f"🔁 Carregando modelo salvo: {args.checkpoint}")
        checkpoint = torch.load(args.checkpoint, map_location=device)
        if "model" in checkpoint:
            model.load_state_dict(checkpoint["model"])
            optimizer.load_state_dict(checkpoint["optimizer"])
            scaler.load_state_dict(checkpoint["scaler"])
            start_epoch = checkpoint.get("epoch", 1) + 1
            best_val_loss = checkpoint.get("best_val_loss", float("inf"))
        else:
            model.load_state_dict(checkpoint)

    # --- Loop de treino ---
    for epoch in range(start_epoch, args.epochs + 1):
        t0 = time.time()

        # ----- Treino -----
        model.train()
        running_train_loss = 0
        for x, y in dl_train:
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad(set_to_none=True)
            with torch.cuda.amp.autocast(enabled=args.amp):
                logits = model(x)
                loss = criterion(logits, y)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            running_train_loss += loss.item() * x.size(0)

        avg_train_loss = running_train_loss / len(ds_train)
        train_loss_history.append(avg_train_loss)

        # ----- Validação + métricas -----
        model.eval()
        running_val_loss = 0
        confmat = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device)

        with torch.no_grad():
            for x, y in dl_val:
                x, y = x.to(device), y.to(device)
                with torch.cuda.amp.autocast(enabled=args.amp):
                    logits = model(x)
                    loss = criterion(logits, y)
                running_val_loss += loss.item() * x.size(0)

                pred = logits.argmax(1)
                confmat = confmat_update(confmat, pred, y, num_classes, ignore_index=ds_train.ignore_id)

        avg_val_loss = running_val_loss / len(ds_val)
        val_loss_history.append(avg_val_loss)

        m = metrics_from_confmat(confmat, main_class_id=main_class_id)
        miou_history.append(m["miou"])
        main_class_f1 = m["main_class"]["f1"] if (m["main_class"] is not None) else None
        if main_class_f1 is not None:
            f1_history.append(main_class_f1)
        cur_lr = optimizer.param_groups[0]["lr"]
        lr_history.append(cur_lr)

        elapsed = time.time() - t0
        msg = (f"[{epoch}/{args.epochs}] "
            f"train_loss={avg_train_loss:.4f} "
            f"val_loss={avg_val_loss:.4f} "
            f"mIoU={m['miou']:.4f} "
            f"pixAcc={m['pixel_acc']:.4f} "
            f"lr={cur_lr:.2e} "
            f"time={elapsed:.1f}s")
        if main_class_f1 is not None:
            msg += f"  | {MAIN_CLASS_NAME}: F1={main_class_f1:.4f} IoU={m['main_class']['iou']:.4f}"
        print(msg)

        # ----- Tracking de melhora por LOSS -----
        improved_loss = avg_val_loss < best_val_loss - 1e-6
        if improved_loss:
            best_val_loss = avg_val_loss
            no_imp_loss = 0
            # checkpoint por loss
            torch.save(model.state_dict(), os.path.join(save_path, f"{MODEL_NAME}_best.pth"))
            torch.save({
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "scaler": scaler.state_dict(),
                "epoch": epoch,
                "best_val_loss": best_val_loss
            }, os.path.join(save_path, f"{MODEL_NAME}_best_checkpoint.pth"))
            print("✅ Novo melhor modelo salvo (val_loss).")
        else:
            no_imp_loss += 1

        # ----- Tracking + checkpoint por F1 da classe principal -----
        if main_class_f1 is not None:
            if main_class_f1 > best_main_class_f1 + delta_f1_min:
                best_main_class_f1 = main_class_f1
                no_imp_f1 = 0
                torch.save(model.state_dict(), os.path.join(save_path, f"{MODEL_NAME}_best_f1_{MAIN_CLASS_NAME}.pth"))
                print(f"🌿💾 Checkpoint salvo (melhor F1 da {MAIN_CLASS_NAME}).")
            else:
                no_imp_f1 += 1
        else:
            # se não houver F1 (ex: id não definido), ignora o critério
            no_imp_f1 = 0

        # ----- Scheduler inteligente -----
        if active_sched == "cosine":
            # se travar por plateau_patience, troca pra ReduceLROnPlateau
            if no_imp_loss >= plateau_patience:
                active_sched = "plateau"
                print("🔁 Mudando scheduler: Cosine → ReduceLROnPlateau (platô detectado).")
                # resets ao trocar
                no_imp_loss = 0
                no_imp_f1   = 0
                epochs_since_switch = 0
                plateau.step(avg_val_loss)  # primeiro passo do plateau
                # (opcional) “adiantar” a queda do LR:
                for g in optimizer.param_groups:
                    g['lr'] = max(g['lr'] * plateau_factor, min_lr)
            else:
                cosine.step()
        else:
            plateau.step(avg_val_loss)
            epochs_since_switch += 1

        # ----- Log de estagnação -----
        print(f"⏳ Sem melhora — loss: {no_imp_loss}/{patience_loss}, {MAIN_CLASS_NAME}: {no_imp_f1}/{patience_f1}")

        # ----- Early stopping bi-critério (com 'graça' após switch) -----
        if (no_imp_loss >= patience_loss and
            (main_class_f1 is None or no_imp_f1 >= patience_f1) and
            (active_sched == "cosine" or epochs_since_switch >= grace_after_switch)):
            print("⏹ Early stopping: loss e F1 sem melhora (com período de graça respeitado).")
            break

        # ----- Plots periódicos -----
        if epoch % 5 == 0 or epoch == args.epochs:
            x_epochs = list(range(start_epoch, start_epoch + len(train_loss_history)))
            # Loss
            plt.figure()
            plt.plot(x_epochs, train_loss_history, marker="o", label="Train Loss")
            plt.plot(x_epochs, val_loss_history, marker="s", label="Val Loss")
            plt.xlabel("Época"); plt.ylabel("Loss"); plt.grid(True); plt.legend(); plt.title("Curva de Loss")
            plt.tight_layout()
            plt.savefig(os.path.join(save_path, "loss_curve.png")); plt.close()
            # LR
            plt.figure()
            plt.plot(x_epochs, lr_history, marker=".")
            plt.xlabel("Época"); plt.ylabel("LR"); plt.grid(True); plt.title("Learning Rate")
            plt.tight_layout()
            plt.savefig(os.path.join(save_path, "lr_curve.png")); plt.close()
            # mIoU e F1(erva)
            plt.figure()
            plt.plot(x_epochs, miou_history, marker="^", label="mIoU")
            if len(f1_history) == len(miou_history):
                plt.plot(x_epochs, f1_history, marker="*", label=f"F1 {MAIN_CLASS_NAME}")
            plt.xlabel("Época"); plt.ylabel("Score"); plt.grid(True); plt.legend(); plt.title(f"mIoU / F1({MAIN_CLASS_NAME})")
            plt.tight_layout()
            plt.savefig(os.path.join(save_path, "metrics_curve.png")); plt.close()

def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--epochs", type=int, default=30)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--amp", action="store_true")
    ap.add_argument("--checkpoint", type=str, default=None, help="Caminho do modelo .pth para continuar o treinamento")
    return ap.parse_args()

if __name__ == "__main__":
    args = parse_args()
    os.makedirs(save_path, exist_ok=True)
    train(args)