# 👉 force backend headless (sem Tk) import itertools import os import random os.environ["MPLBACKEND"] = "Agg" # extra-garantia import matplotlib matplotlib.use("Agg") # tem que vir antes do pyplot! import matplotlib.pyplot as plt plt.ioff() # desliga modo interativo import json import time import argparse from PIL import Image import numpy as np import torch import torch.nn.functional as F 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 from utils import carregar_labelmap_completo, compute_roi_indices # ⚙️ 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 = 32 num_workers = 4 # ---- Helpers de métricas ---- def _infer_ignore_id(ignore_rgb, default_id=255): import numpy as _np if isinstance(ignore_rgb, (list, tuple)): if len(ignore_rgb) == 1 and isinstance(ignore_rgb[0], (int, _np.integer)): return int(ignore_rgb[0]) if len(ignore_rgb) == 3: return default_id if isinstance(ignore_rgb, (int, _np.integer)): return int(ignore_rgb) return default_id def compute_class_weights_from_split( train_split_root: str, labelmap_path: str, roi_inicio: float, roi_tamanho: float, *, alpha: float = 1.2, # ↑ reforça classes raras (1.0 a 1.5 costuma ir bem) w_min: float = 0.3, # piso geral w_max: float = 4.0, # teto geral floor_bg: float = 0.4, # piso específico pro 'chao' / background max_samples_per_group: int = 300 # amostras p/ grupo (acelera o cálculo) ): """ Calcula pesos dinâmicos (median frequency balancing ^ alpha) SOBRE A ROI das máscaras do split/train. - Normaliza por ROI (mesma fatia usada no dataset). - Clampa pesos entre [w_min, w_max] e mantém 'chao' no mínimo floor_bg. Retorna: tensor de pesos (indexado por ID de classe). """ # Labelmap _, _, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path) ignore_id = _infer_ignore_id(ignore_rgb, default_id=255) class_ids = sorted(classes.keys()) # ex.: [0,1,2] counts = np.zeros(len(class_ids), dtype=np.int64) # Onde estão as máscaras do train? group_root = os.path.join(train_split_root, "group") group_dirs = [] if os.path.isdir(group_root): for g in sorted(os.listdir(group_root)): mdir = os.path.join(group_root, g, "masks") if os.path.isdir(mdir): group_dirs.append(mdir) else: # fallback legado mdir = os.path.join(train_split_root, "masks") if os.path.isdir(mdir): group_dirs.append(mdir) # Conta pixels por classe DENTRO DA ROI for mdir in group_dirs: n = 0 for p in os.listdir(mdir): if not p.lower().endswith(".png"): continue m = np.array(Image.open(os.path.join(mdir, p)).convert("L")) H = m.shape[0] y_fim, y_ini = compute_roi_indices(H, roi_inicio, roi_tamanho) roi = m[y_fim:y_ini, :] # ignora 'ignore' e só soma classes válidas for i, cid in enumerate(class_ids): if cid == ignore_id: continue counts[i] += int((roi == cid).sum()) n += 1 if n >= max_samples_per_group: break total = int(counts.sum()) if total == 0: # fallback seguro print("⚠️ compute_class_weights_from_split: não encontrei pixels válidos; usando pesos [1,1,...].") return None, None freqs = counts / total # frequência por classe nonzero = freqs[freqs > 0] base = np.median(nonzero) if nonzero.size > 0 else 1.0 weights_arr = np.zeros_like(freqs, dtype=np.float32) for i, f in enumerate(freqs): if f <= 0: w = w_max else: # median-freq ^ alpha w = (base / f) ** alpha w = float(np.clip(w, w_min, w_max)) weights_arr[i] = w # Piso do 'chao' (ou 'background'), se existir for i, cid in enumerate(class_ids): name = str(classes[cid]).lower() if ("chao" in name) or ("background" in name): weights_arr[i] = max(weights_arr[i], floor_bg) # Constrói vetor na indexação por ID de classe (0..max_id) max_cid = max(class_ids) weights_full = np.ones(max_cid + 1, dtype=np.float32) for i, cid in enumerate(class_ids): weights_full[cid] = weights_arr[i] # Log bonitinho pretty = {int(cid): (str(classes[cid]), float(weights_full[cid]), float(freqs[i])) for i, cid in enumerate(class_ids)} return weights_full, pretty @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 _id_by_name(d, name): name = name.lower() for cid, nm in d.items(): if isinstance(nm, str) and name in nm.lower(): return cid return None def _get_mask_roi_from_ds(ds, i, roi_inicio, roi_tamanho): """Tenta obter o caminho da máscara; se não der, usa ds[i].""" mask_path = None if hasattr(ds, "mask_paths"): mask_path = ds.mask_paths[i] elif hasattr(ds, "items"): item = ds.items[i] if isinstance(item, dict) and "mask" in item: mask_path = item["mask"] if mask_path is not None: m = np.array(Image.open(mask_path).convert("L")) else: # fallback: carrega a máscara já processada pelo dataset _, y = ds[i] # y: Tensor [H,W] m = y.cpu().numpy() H = m.shape[0] y_fim, y_ini = compute_roi_indices(H, roi_inicio, roi_tamanho) return m[y_fim:y_ini, :] def compute_presence_indices(ds, class_ids, roi_inicio, roi_tamanho): """ presence: {cid: [idxs que CONTÊM essa classe na ROI]} others: [idxs que NÃO contêm NENHUMA das 'class_ids' na ROI] """ presence = {cid: [] for cid in class_ids} others = [] for i in range(len(ds)): roi = _get_mask_roi_from_ds(ds, i, roi_inicio, roi_tamanho) found_any = False for cid in class_ids: if (roi == cid).any(): presence[cid].append(i) found_any = True if not found_any: others.append(i) return presence, others class EnsureClassesBatchSampler(torch.utils.data.Sampler): """ Garante >=1 amostra de CADA classe em 'required_classes' por batch. Preenche o resto com índices do pool (others + todo o conjunto). Use com DataLoader(..., batch_sampler= sampler) sem passar batch_size/sampler/shuffle. """ def __init__(self, presence, total_indices, batch_size, required_classes, seed=42): self.presence = presence # dict cid -> list[idx] self.required = [c for c in required_classes if len(presence.get(c, [])) > 0] self.batch_size = batch_size # iteradores cíclicos (com reposição) por classe requerida self.iters = { c: itertools.cycle(self.presence[c]) for c in self.required } # pool de preenchimento: todos os índices (mistura bem) self.rest_iter = itertools.cycle(list(total_indices)) self.rng = random.Random(seed) # tamanho lógico: nº de batches por época self._length = max(1, int(np.ceil(len(total_indices) / float(batch_size)))) def __iter__(self): for _ in range(self._length): batch = [] # 1 de cada classe requerida (se existir) for c in self.required: batch.append(next(self.iters[c])) # completa o batch while len(batch) < self.batch_size: batch.append(next(self.rest_iter)) self.rng.shuffle(batch) yield batch def __len__(self): return self._length 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 ) # --- Mapa id->nome (usa o do dataset; se não houver, carrega do labelmap) --- if hasattr(ds_train, "classes") and isinstance(ds_train.classes, dict) and len(ds_train.classes) > 0: id_to_name = {int(k): str(v) for k, v in ds_train.classes.items()} else: # fallback seguro ao arquivo de labelmap _, _, id_to_name, _ = carregar_labelmap_completo(labelmap_path) id_to_name = {int(k): str(v) for k, v in id_to_name.items()} # name->id (case-insensitive) name_to_id = {v.strip().lower(): k for k, v in id_to_name.items()} # --- Lista dinâmica de classes a garantir por batch --- raw = getattr(args, "ensure_per_batch", "") req_names = [s.strip().lower() for s in raw.split(",") if s.strip()] req_ids = [] for nm in req_names: cid = name_to_id.get(nm) if cid is None: # tenta correspondência parcial (p.ex. "erva" casa com "Erva", "weed_erva", etc.) matches = [k for k, v in id_to_name.items() if nm in v.lower()] if len(matches) == 1: cid = matches[0] elif len(matches) > 1: print(f"⚠️ '--ensure-per-batch {nm}': ambíguo entre {[id_to_name[m] for m in matches]}; ignorando este nome.") cid = None else: print(f"⚠️ '--ensure-per-batch {nm}': classe não encontrada nas classes {list(name_to_id.keys())}.") if cid is not None and cid not in req_ids: req_ids.append(cid) if len(req_ids) > 0: presence, _ = compute_presence_indices( ds_train, class_ids=req_ids, roi_inicio=ROI_INICIO, roi_tamanho=ROI_TAMANHO ) total_indices = range(len(ds_train)) batch_sampler = EnsureClassesBatchSampler( presence=presence, total_indices=total_indices, batch_size=batch_size, required_classes=req_ids, seed=getattr(args, "seed", 42) ) # ⚠️ Use 'batch_sampler' (NÃO passe batch_size/sampler/shuffle) dl_train = DataLoader(ds_train, batch_sampler=batch_sampler, num_workers=num_workers, pin_memory=True) else: dl_train = DataLoader(ds_train, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True) # Validação normal 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) # --- Pesos dinâmicos por classe (sobre a ROI do split/train) --- train_split_root = os.path.join(dataset_path, "split", "train") weights_np, debug_info = compute_class_weights_from_split( train_split_root=train_split_root, labelmap_path=labelmap_path, roi_inicio=ROI_INICIO, roi_tamanho=ROI_TAMANHO, alpha=getattr(args, "cw_alpha", 1.05), w_min=getattr(args, "cw_min", 0.3), w_max=getattr(args, "cw_max", 2.0), floor_bg=getattr(args, "cw_bgfloor", 0.6), max_samples_per_group=getattr(args, "cw_max_per_group", 300) ) if weights_np is None or not getattr(args, "use_wights", False): # fallback seguro weights_t = None print("⚠️ Pesos dinâmicos indisponíveis; usando CrossEntropy sem pesos.") else: import pprint pprint.pprint({"class_weights": debug_info}) weights_t = torch.tensor(weights_np, device=device) # --- Warm-up de pesos por época --- ones = torch.ones_like(weights_t) if weights_t is not None else None def make_epoch_weights(epoch, *, cw_warmup=8): """ Interpola: w_epoch = (1 - λ) * 1 + λ * weights_t λ cresce de 0→1 nas primeiras `cw_warmup` épocas. Retorna (w_epoch, dice_w_normalized) ou (None, None) se sem pesos. """ if weights_t is None: return None, None # lê da CLI ou usa default cw_warmup = getattr(args, "cw_warmup", cw_warmup) # λ linear 0→1 (pode trocar por cosseno, ver abaixo) cw_lambda = min(1.0, max(0.0, (epoch - 1) / max(1, cw_warmup))) w_epoch = (1.0 - cw_lambda) * ones + cw_lambda * weights_t # normaliza para o Dice (evita distorção) dice_w = (w_epoch / w_epoch.mean()).detach() return w_epoch, dice_w optimizer = optim.AdamW(model.parameters(), lr=getattr(args, "lr", 3e-4), weight_decay=1e-4) def dice_loss(logits, target, ignore_index=255, class_weights=None, eps=1e-6): """ logits: [N, C, H, W] (antes do softmax) target: [N, H, W] com IDs de classe; 'ignore_index' será mascarado class_weights: tensora opcional [C] (ex.: pesos da CE, normalizados) """ N, C, H, W = logits.shape # Probabilidades por classe pred = F.softmax(logits, dim=1) # [N,C,H,W] # Máscara de válidos (ignora 255) valid = (target != ignore_index) # [N,H,W] target_clamped = torch.clamp(target, 0, C-1) # evita index out of range # One-hot do target (com válidos) one_hot = torch.zeros((N, C, H, W), device=logits.device, dtype=pred.dtype) one_hot.scatter_(1, target_clamped.unsqueeze(1), 1.0) # [N,1,H,W] -> [N,C,H,W] # Aplica máscara de válidos valid = valid.unsqueeze(1) # [N,1,H,W] pred = pred * valid one_hot = one_hot * valid # Dice por classe (agrega em N,H,W) inter = (pred * one_hot).sum(dim=(0, 2, 3)) # [C] pred_sum = pred.sum(dim=(0, 2, 3)) # [C] tgt_sum = one_hot.sum(dim=(0, 2, 3)) # [C] dice = (2 * inter + eps) / (pred_sum + tgt_sum + eps) # [C] if class_weights is not None: # opcional: ponderar o Dice com pesos (normalize antes!) # garante shape [C] w = torch.ones(C, device=logits.device, dtype=pred.dtype) w[:class_weights.numel()] = class_weights loss = 1.0 - (w * dice).sum() / (w.sum() + eps) else: loss = 1.0 - dice.mean() return loss # 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.amp.GradScaler('cuda', 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() # pesos deste epoch w_epoch, dice_w = make_epoch_weights(epoch) dice_mix = min(0.3, (epoch-1)/10 * 0.3) # 0.0→0.3 nas 10 primeiras ce_mix = 1.0 - dice_mix if w_epoch is not None: erva_id = _id_by_name(ds_train.classes, "erva") cana_id = _id_by_name(ds_train.classes, "cana") chao_id = _id_by_name(ds_train.classes, "chao") # calcula lambda atual (igual ao make_epoch_weights) cw_warmup = getattr(args, "cw_warmup", 8) cw_lambda = min(1.0, max(0.0, (epoch - 1) / max(1, cw_warmup))) if cw_lambda < 1.0: if erva_id is not None: w_epoch[erva_id] = torch.clamp(w_epoch[erva_id], min=1.2) if cana_id is not None: w_epoch[cana_id] = torch.clamp(w_epoch[cana_id], max=2.2) if chao_id is not None: w_epoch[chao_id] = torch.clamp(w_epoch[chao_id], min=0.5) # re-normaliza o peso do Dice após clamps dice_w = (w_epoch / w_epoch.mean()).detach() ce = nn.CrossEntropyLoss(ignore_index=ds_train.ignore_id, weight=w_epoch) else: ce = nn.CrossEntropyLoss(ignore_index=ds_train.ignore_id) # ----- 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.amp.autocast('cuda', enabled=args.amp): logits = model(x) dloss = dice_loss(logits, y, ignore_index=ds_train.ignore_id, class_weights=None) loss = ce_mix * ce(logits, y) + dice_mix * dloss 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.amp.autocast('cuda', enabled=args.amp): logits = model(x) dloss = dice_loss(logits, y, ignore_index=ds_train.ignore_id, class_weights=None) loss = ce_mix * ce(logits, y) + dice_mix * dloss 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 ----- if no_imp_loss > 0 or no_imp_f1 > 0: 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") ap.add_argument("--ensure-per-batch", type=str, default="", help="Lista de classes por nome para garantir >=1 por batch. Ex.: 'erva,cana'") ap.add_argument("--use-weights", action="store_true") return ap.parse_args() if __name__ == "__main__": args = parse_args() os.makedirs(save_path, exist_ok=True) train(args)