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)