#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ _8_train_segformer_oak_fcc3_multihead.py Treina SegFormer OAK-FCC-3 com tensor multispectral [R,G,B,RE,NIR] e contrato multi-head: 1) semantic_head: máscara: masks/.npy classes: 0=chao, 1=cana, 2=erva, 255=ignore 2) vegetation_head: máscara: masks_vegetation/.npy classes: 0=nao_vegetacao, 1=vegetacao, 255=ignore 3) cana_head: máscara: masks_cana/.npy classes: 0=nao_cana, 1=cana, 255=ignore Decisão operacional futura: alvo/pulverizavel = vegetation == 1 AND cana == 0 Entrada esperada após normalize + split: dataset/split/train/group// tensors/.npy masks/.npy masks_vegetation/.npy masks_cana/.npy metas/.json previews/.png dataset/split/val/group//... Exemplo: python _8_train_segformer_oak_fcc3_multihead.py ^ --epochs 50 --batch 2 --lr 3e-5 --wd 0.01 ^ --num_workers 2 --amp --amp_val --grad_accum 4 ^ --class_weights auto """ from __future__ import annotations import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128") import csv import copy import json import time import argparse import random from pathlib import Path from typing import Dict, List, Optional, Tuple, Any import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from torch.amp import autocast, GradScaler from transformers import SegformerForSemanticSegmentation # ============================================================ # Seed / util # ============================================================ def set_seed(seed: int = 42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def ensure_dir(path: Path): path.mkdir(parents=True, exist_ok=True) def load_json(path: str | Path) -> dict: with open(path, "r", encoding="utf-8") as f: return json.load(f) def save_json(path: Path, data: dict): ensure_dir(path.parent) with path.open("w", encoding="utf-8") as f: json.dump(data, f, ensure_ascii=False, indent=2) def safe_float(x, default=0.0) -> float: try: return float(x) except Exception: return float(default) # ============================================================ # Labelmap # ============================================================ def load_labelmap(labelmap_path: str): """ Tenta usar helpers.carregar_labelmap_completo. Fallback simples para labelmap com uma classe por linha. """ try: from helpers import carregar_labelmap_completo, _infer_ignore_id _cor_para_id, _colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path) ignore_id = _infer_ignore_id(ignore_rgb, 255) id2label = {int(k): str(v) for k, v in id_para_nome.items()} label2id = {v.lower(): k for k, v in id2label.items()} return id2label, label2id, int(ignore_id) except Exception as e: print(f"[WARN] Não consegui usar helpers.carregar_labelmap_completo: {e}") print("[WARN] Usando parser simples: uma classe por linha.") id2label = {} with open(labelmap_path, "r", encoding="utf-8") as f: for line in f: s = line.strip() if not s or s.startswith("#"): continue parts = s.replace(",", " ").split() if len(parts) >= 2 and parts[0].isdigit(): cid = int(parts[0]) name = parts[1] else: cid = len(id2label) name = parts[0] if name.lower() in ("ignore", "void", "background_ignore"): continue id2label[cid] = name label2id = {v.lower(): k for k, v in id2label.items()} return id2label, label2id, 255 # ============================================================ # Head config # ============================================================ DEFAULT_HEADS = { "semantic": { "enabled": True, "type": "multiclass", "num_classes": 3, "mask_dir": "masks", "classes": { "chao": 0, "cana": 1, "erva": 2, }, "ignore_index": 255, "loss_weight": 0.10, }, "vegetation": { "enabled": True, "type": "binary", "num_classes": 2, "mask_dir": "masks_vegetation", "classes": { "background": 0, "vegetation": 1, }, "ignore_index": 255, "loss_weight": 0.20, }, "cana": { "enabled": True, "type": "binary", "num_classes": 2, "mask_dir": "masks_cana", "classes": { "not_cana": 0, "cana": 1, }, "ignore_index": 255, "loss_weight": 0.25, }, "target": { "enabled": True, "type": "binary", "num_classes": 2, "mask_dir": "__derived_target__", "classes": { "background": 0, "target": 1, }, "ignore_index": 255, "loss_weight": 0.45, "derived_from": ["vegetation", "cana"], }, } def merge_dict(dst: dict, src: dict) -> dict: out = copy.deepcopy(dst) def rec(a, b): for k, v in b.items(): if isinstance(v, dict) and isinstance(a.get(k), dict): rec(a[k], v) else: a[k] = v if isinstance(src, dict): rec(out, src) return out def build_heads_config(config: dict, ignore_index: int) -> Dict[str, dict]: cfg = merge_dict(DEFAULT_HEADS, config.get("heads", {}) or {}) # Compatibilidade com configs antigas. if not bool(config.get("multi_head", True)): print("[WARN] config.multi_head não está true. Este script vai treinar multi-head mesmo assim.") for name, hcfg in cfg.items(): hcfg.setdefault("enabled", True) hcfg.setdefault("ignore_index", ignore_index) hcfg["num_classes"] = int(hcfg.get("num_classes", 2)) hcfg["loss_weight"] = float(hcfg.get("loss_weight", 1.0)) hcfg["mask_dir"] = str(hcfg.get("mask_dir", "masks")) active = {k: v for k, v in cfg.items() if bool(v.get("enabled", True))} if "semantic" not in active: raise RuntimeError("A head 'semantic' deve estar habilitada neste contrato.") if "vegetation" not in active: raise RuntimeError("A head 'vegetation' deve estar habilitada neste contrato.") if "cana" not in active: raise RuntimeError("A head 'cana' deve estar habilitada neste contrato.") total_w = sum(float(v.get("loss_weight", 0.0)) for v in active.values()) if total_w <= 0: raise RuntimeError("Soma de loss_weight das heads deve ser > 0.") # Normaliza os pesos para não bagunçar magnitude da loss. for v in active.values(): v["loss_weight_norm"] = float(v.get("loss_weight", 0.0)) / total_w return active # ============================================================ # Dataset OAK-FCC-3 multi-head # ============================================================ class OakFcc3TensorMultiHeadDataset(Dataset): """ Lê o contrato pós-normalização/split: root/ group//tensors/.npy group//masks/.npy group//masks_vegetation/.npy group//masks_cana/.npy """ def __init__( self, root: str | Path, heads_config: Dict[str, dict], channels: int = 5, channel_indices: Optional[List[int]] = None, strict_channels: bool = False, resize_hw: Optional[Tuple[int, int]] = None, ): self.root = Path(root) self.heads_config = heads_config self.channels = int(channels) self.strict_channels = bool(strict_channels) self.resize_hw = resize_hw self.channel_indices = channel_indices self.samples = self._collect_samples() if not self.samples: raise RuntimeError(f"Nenhuma amostra encontrada em: {self.root}") def _collect_samples(self): samples = [] group_root = self.root / "group" if not group_root.is_dir(): raise RuntimeError(f"Pasta group não encontrada em: {self.root}") for group_dir in sorted(group_root.iterdir()): if not group_dir.is_dir(): continue tensors_dir = group_dir / "tensors" metas_dir = group_dir / "metas" previews_dir = group_dir / "previews" if not tensors_dir.is_dir(): continue for tensor_path in sorted(tensors_dir.glob("*.npy")): base = tensor_path.stem masks = {} missing = [] for head_name, hcfg in self.heads_config.items(): mask_dir_name = str(hcfg.get("mask_dir")) # Head target derivada: não precisa arquivo .npy físico if mask_dir_name == "__derived_target__" or bool(hcfg.get("derived", False)): masks[head_name] = None continue mask_dir = group_dir / mask_dir_name mask_path = mask_dir / f"{base}.npy" if not mask_path.exists(): missing.append(f"{head_name}:{mask_path}") else: masks[head_name] = mask_path if missing: print(f"[WARN] Pulando {tensor_path}, masks ausentes: {missing}") continue meta_path = metas_dir / f"{base}.json" preview_path = previews_dir / f"{base}.png" samples.append({ "group": group_dir.name, "base": base, "tensor": tensor_path, "masks": masks, "meta": meta_path if meta_path.exists() else None, "preview": preview_path if preview_path.exists() else None, }) return samples def __len__(self): return len(self.samples) def _resize_tensor(self, x: torch.Tensor): if self.resize_hw is None: return x h, w = self.resize_hw if x.shape[-2:] != (h, w): x = F.interpolate( x.unsqueeze(0), size=(h, w), mode="bilinear", align_corners=False, ).squeeze(0) return x def _resize_mask(self, y: torch.Tensor): if self.resize_hw is None: return y h, w = self.resize_hw if y.shape[-2:] != (h, w): y = F.interpolate( y.unsqueeze(0).unsqueeze(0).float(), size=(h, w), mode="nearest", ).squeeze(0).squeeze(0).long() return y def __getitem__(self, idx): s = self.samples[idx] x = np.load(str(s["tensor"])).astype(np.float32) if x.ndim != 3: raise RuntimeError(f"Tensor inválido {s['tensor']}: shape={x.shape}") if self.channel_indices is not None: max_idx = max(self.channel_indices) if x.shape[0] <= max_idx: raise RuntimeError( f"Tensor {s['tensor']} tem {x.shape[0]} canais, " f"mas precisa acessar índice {max_idx}. Shape={x.shape}" ) x = x[self.channel_indices, :, :] else: if x.shape[0] < self.channels: raise RuntimeError( f"Tensor {s['tensor']} tem {x.shape[0]} canais, " f"mas config pediu {self.channels}." ) x = x[:self.channels] xt = torch.from_numpy(np.ascontiguousarray(x)).float() xt = self._resize_tensor(xt) masks = {} # Primeiro carrega as heads com arquivo físico for head_name, path in s["masks"].items(): hcfg = self.heads_config[head_name] if path is None: continue y = np.load(str(path)).astype(np.int64) yt = torch.from_numpy(np.ascontiguousarray(y)).long() yt = self._resize_mask(yt) masks[head_name] = yt # Depois deriva a target, se existir no contrato if "target" in self.heads_config: if "vegetation" not in masks or "cana" not in masks: raise RuntimeError("Head target requer masks vegetation e cana carregadas.") ignore_index = int(self.heads_config["target"].get("ignore_index", 255)) veg = masks["vegetation"] cana = masks["cana"] valid = (veg != ignore_index) & (cana != ignore_index) target = torch.zeros_like(veg, dtype=torch.long) target[(veg == 1) & (cana == 0)] = 1 target[~valid] = ignore_index masks["target"] = target return { "image": xt, "masks": masks, "group": s["group"], "base": s["base"], } def collate_fn(batch): imgs = torch.stack([b["image"] for b in batch], dim=0) head_names = list(batch[0]["masks"].keys()) masks = { h: torch.stack([b["masks"][h] for b in batch], dim=0) for h in head_names } meta = { "group": [b["group"] for b in batch], "base": [b["base"] for b in batch], } return imgs, masks, meta # ============================================================ # Normalização # ============================================================ class FixedNormalizer(nn.Module): def __init__(self, mean: List[float], std: List[float]): super().__init__() mean_t = torch.tensor(mean, dtype=torch.float32).view(1, -1, 1, 1) std_t = torch.tensor(std, dtype=torch.float32).view(1, -1, 1, 1) self.register_buffer("mean", mean_t) self.register_buffer("std", torch.clamp(std_t, min=1e-6)) def forward(self, x): return (x - self.mean) / self.std def normalize_per_batch(x: torch.Tensor, eps: float = 1e-6): mean = x.mean(dim=(0, 2, 3), keepdim=True) std = x.std(dim=(0, 2, 3), keepdim=True).clamp_min(eps) return (x - mean) / std def build_normalizer(config: dict, args, device: torch.device): input_channel_names = get_input_channel_names(config) input_channel_indices = get_input_channel_indices(config) channels = len(input_channel_names) model_family = config.get("modelo", "segformer") model_name = config.get("model_name", "model") stats_source_tag = config.get("stats_source_tag", f"stacked_raw{channels}") candidates = [] if args.norm_stats: candidates.append(Path(args.norm_stats)) candidates.append(Path("dataset") / f"{config['resolucao'][0]}x{config['resolucao'][1]}" / "group" / "norm_stats.json") candidates.append(Path("backup") / model_family / model_name / stats_source_tag / "norm_stats.json") for p in candidates: if p.exists(): stats = load_json(p) mean = stats.get("mean", []) std = stats.get("std", []) stat_channels = stats.get("channels", []) if len(mean) < max(input_channel_indices) + 1 or len(std) < max(input_channel_indices) + 1: raise RuntimeError( f"norm_stats incompatível: {p} | " f"precisa índices={input_channel_indices}, " f"mean={len(mean)}, std={len(std)}" ) mean = [mean[i] for i in input_channel_indices] std = [std[i] for i in input_channel_indices] if stat_channels: stat_channels = [stat_channels[i] for i in input_channel_indices] print(f"[NORM] usando stats fixos: {p}") print(f"[NORM] selected_channels={stat_channels or input_channel_names}") return FixedNormalizer(mean, std).to(device), str(p) print("[NORM] norm_stats não encontrado. Usando normalize_per_batch.") return None, None DEFAULT_CHANNEL_ORDER = ["R", "G", "B", "RE", "NIR"] def get_input_channel_names(config: dict) -> List[str]: """ Define quais canais entram no modelo. Padrão: channels=3 -> R,G,B channels=4 -> R,G,B,RE channels=5 -> R,G,B,RE,NIR Futuramente permite: "input_channels": ["R", "G", "B", "NIR"] """ if "input_channels" in config: names = [str(c).upper() for c in config["input_channels"]] else: n = int(config.get("channels", 5)) names = DEFAULT_CHANNEL_ORDER[:n] invalid = [c for c in names if c not in DEFAULT_CHANNEL_ORDER] if invalid: raise RuntimeError(f"Canais inválidos em input_channels: {invalid}") return names def get_input_channel_indices(config: dict) -> List[int]: names = get_input_channel_names(config) return [DEFAULT_CHANNEL_ORDER.index(c) for c in names] # ============================================================ # Modelo multi-head # ============================================================ def patch_segformer_encoder_input_channels(segformer_encoder: nn.Module, in_ch: int): """ Altera o primeiro patch embedding do SegFormer para aceitar C canais. Inicializa canais extras pela média dos pesos RGB. """ if in_ch == 3: return segformer_encoder proj = segformer_encoder.encoder.patch_embeddings[0].proj if proj.in_channels == in_ch: return segformer_encoder old_weight = proj.weight.data.clone() old_bias = proj.bias.data.clone() if proj.bias is not None else None new_proj = nn.Conv2d( in_channels=in_ch, out_channels=proj.out_channels, kernel_size=proj.kernel_size, stride=proj.stride, padding=proj.padding, dilation=proj.dilation, groups=proj.groups, bias=proj.bias is not None, padding_mode=proj.padding_mode, ) with torch.no_grad(): if in_ch <= old_weight.shape[1]: new_proj.weight.copy_(old_weight[:, :in_ch, :, :]) else: new_proj.weight[:, :old_weight.shape[1], :, :].copy_(old_weight) extra = in_ch - old_weight.shape[1] mean_w = old_weight.mean(dim=1, keepdim=True) new_proj.weight[:, old_weight.shape[1]:, :, :].copy_(mean_w.repeat(1, extra, 1, 1)) if old_bias is not None: new_proj.bias.copy_(old_bias) segformer_encoder.encoder.patch_embeddings[0].proj = new_proj print(f"[MODEL] patch input channels: 3 -> {in_ch}") return segformer_encoder def replace_segformer_decode_classifier(decode_head: nn.Module, num_classes: int): """ Troca o classifier final da decode_head do HuggingFace SegFormer. """ if not hasattr(decode_head, "classifier"): raise RuntimeError("decode_head sem atributo classifier. Estrutura SegFormer inesperada.") old = decode_head.classifier if not isinstance(old, nn.Conv2d): raise RuntimeError(f"decode_head.classifier não é Conv2d: {type(old)}") new = nn.Conv2d( in_channels=old.in_channels, out_channels=int(num_classes), kernel_size=old.kernel_size, stride=old.stride, padding=old.padding, dilation=old.dilation, groups=old.groups, bias=old.bias is not None, padding_mode=old.padding_mode, ) nn.init.xavier_uniform_(new.weight) if new.bias is not None: nn.init.zeros_(new.bias) decode_head.classifier = new return decode_head class MultiHeadSegFormer(nn.Module): def __init__( self, backbone: str, channels: int, heads_config: Dict[str, dict], semantic_id2label: Dict[int, str], semantic_label2id: Dict[str, int], ): super().__init__() semantic_classes = int(heads_config["semantic"].get("num_classes", len(semantic_id2label))) base = SegformerForSemanticSegmentation.from_pretrained( backbone, num_labels=semantic_classes, id2label={int(k): str(v) for k, v in semantic_id2label.items()}, label2id={str(k): int(v) for k, v in semantic_label2id.items()}, ignore_mismatched_sizes=True, ) patch_segformer_encoder_input_channels(base.segformer, channels) base.config.num_channels = int(channels) self.segformer = base.segformer self.heads_config = heads_config self.decode_heads = nn.ModuleDict() for head_name, hcfg in heads_config.items(): h = copy.deepcopy(base.decode_head) h = replace_segformer_decode_classifier(h, int(hcfg["num_classes"])) self.decode_heads[head_name] = h self.config = base.config def forward(self, pixel_values: torch.Tensor) -> Dict[str, torch.Tensor]: outputs = self.segformer( pixel_values=pixel_values, output_hidden_states=True, return_dict=True, ) hidden_states = outputs.hidden_states logits = {} for head_name, head in self.decode_heads.items(): logits[head_name] = head(hidden_states) return logits def build_model( backbone: str, channels: int, heads_config: Dict[str, dict], semantic_id2label: Dict[int, str], semantic_label2id: Dict[str, int], ): model = MultiHeadSegFormer( backbone=backbone, channels=channels, heads_config=heads_config, semantic_id2label=semantic_id2label, semantic_label2id=semantic_label2id, ) return model # ============================================================ # Métricas / loss # ============================================================ @torch.no_grad() def update_confusion_matrix(cm, preds, labels, num_classes, ignore_index=255): preds = preds.reshape(-1) labels = labels.reshape(-1) valid = labels != ignore_index preds = preds[valid] labels = labels[valid] valid2 = (labels >= 0) & (labels < num_classes) & (preds >= 0) & (preds < num_classes) preds = preds[valid2] labels = labels[valid2] if labels.numel() == 0: return 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, eps=1e-6): 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, eps=1e-6): cm = cm.float() return float(torch.diag(cm).sum() / (cm.sum() + eps)) def dice_loss(logits, target, num_classes, ignore_index=255, smooth=1.0): probs = torch.softmax(logits, dim=1) valid = target != ignore_index if valid.sum() == 0: return logits.new_tensor(0.0) target_clamped = target.clone() target_clamped[~valid] = 0 target_clamped = target_clamped.long() target_1h = F.one_hot(target_clamped, num_classes=num_classes) target_1h = target_1h.permute(0, 3, 1, 2).float() valid_f = valid.unsqueeze(1).float() probs = probs * valid_f target_1h = target_1h * valid_f dims = (0, 2, 3) inter = (probs * target_1h).sum(dims) den = probs.sum(dims) + target_1h.sum(dims) dice = (2.0 * inter + smooth) / (den + smooth) return 1.0 - dice.mean() def build_target_teacher_from_heads( logits_by_head: Dict[str, torch.Tensor], target_shape: Tuple[int, int], config: dict, erva_id: int = 2, ): cfg = config.get("target_distillation", {}) or {} if not bool(cfg.get("enabled", False)): return None required = ("semantic", "vegetation", "cana") if any(k not in logits_by_head for k in required): return None sem = logits_by_head["semantic"] veg = logits_by_head["vegetation"] cana = logits_by_head["cana"] if sem.shape[-2:] != target_shape: sem = F.interpolate(sem, size=target_shape, mode="bilinear", align_corners=False) if veg.shape[-2:] != target_shape: veg = F.interpolate(veg, size=target_shape, mode="bilinear", align_corners=False) if cana.shape[-2:] != target_shape: cana = F.interpolate(cana, size=target_shape, mode="bilinear", align_corners=False) p_sem = torch.softmax(sem, dim=1) p_veg = torch.softmax(veg, dim=1) p_cana = torch.softmax(cana, dim=1) p_sem_erva = p_sem[:, int(erva_id), :, :] p_veg_pos = p_veg[:, 1, :, :] p_cana_pos = p_cana[:, 1, :, :] w_sem_erva = float(cfg.get("w_sem_erva", 0.45)) w_veg_not_cana = float(cfg.get("w_veg_not_cana", 0.35)) w_veg_suppressed = float(cfg.get("w_veg_suppressed", 0.20)) power = float(cfg.get("cana_suppression_power", 1.5)) not_cana = torch.clamp(1.0 - p_cana_pos, 0.0, 1.0) teacher = ( w_sem_erva * p_sem_erva + w_veg_not_cana * p_veg_pos * not_cana + w_veg_suppressed * p_veg_pos * torch.pow(not_cana, power) ) teacher_min = float(cfg.get("teacher_min", 0.0)) teacher_max = float(cfg.get("teacher_max", 1.0)) teacher = torch.clamp(teacher, teacher_min, teacher_max) if bool(cfg.get("detach_teacher", True)): teacher = teacher.detach() return teacher def estimate_head_class_weights( ds: Dataset, head_name: str, num_classes: int, ignore_index: int = 255, max_samples: int = 800, seed: int = 42, ): rng = np.random.default_rng(seed) n = min(len(ds), max_samples) idxs = rng.choice(len(ds), size=n, replace=False) counts = np.zeros(num_classes, dtype=np.float64) for i in idxs: item = ds[i] m = item["masks"][head_name].numpy().reshape(-1) m = m[m != ignore_index] m = m[(m >= 0) & (m < num_classes)] if m.size > 0: 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), counts def build_criterions( ds_train: Dataset, heads_config: Dict[str, dict], class_weights_mode: str, device: torch.device, seed: int, ): criterions = {} weight_debug = {} for head_name, hcfg in heads_config.items(): num_classes = int(hcfg["num_classes"]) ignore_index = int(hcfg.get("ignore_index", 255)) mode = str(class_weights_mode).lower() weights = None counts = None if mode == "none": weights = None elif mode == "auto": w, counts = estimate_head_class_weights( ds_train, head_name=head_name, num_classes=num_classes, ignore_index=ignore_index, seed=seed, ) weights = w.to(device) print(f"[LOSS:{head_name}] class counts:", counts.astype(int).tolist()) print(f"[LOSS:{head_name}] class weights:", w.cpu().numpy().round(3).tolist()) else: # Formato opcional: # semantic=1,2,3;vegetation=1,2;cana=1,4 # ou, para compatibilidade, uma lista aplicada só na semantic. parsed = parse_class_weights_string(class_weights_mode) if head_name in parsed: parts = parsed[head_name] elif head_name == "semantic" and "__single__" in parsed: parts = parsed["__single__"] else: parts = None if parts is not None: if len(parts) != num_classes: raise RuntimeError( f"Pesos da head {head_name} precisam ter {num_classes} valores. Veio: {parts}" ) weights = torch.tensor(parts, dtype=torch.float32, device=device) criterions[head_name] = nn.CrossEntropyLoss( weight=weights, ignore_index=ignore_index, ) weight_debug[head_name] = { "weights": weights.detach().cpu().tolist() if weights is not None else None, "counts": counts.astype(int).tolist() if counts is not None else None, } return criterions, weight_debug def parse_class_weights_string(s: str) -> Dict[str, List[float]]: out = {} txt = str(s).strip() if not txt: return out if "=" not in txt: out["__single__"] = [float(x) for x in txt.split(",") if x.strip()] return out for block in txt.split(";"): block = block.strip() if not block: continue k, v = block.split("=", 1) out[k.strip()] = [float(x) for x in v.split(",") if x.strip()] return out @torch.no_grad() def update_operational_target_cm( cm_target, pred_veg: torch.Tensor, pred_cana: torch.Tensor, gt_veg: torch.Tensor, gt_cana: torch.Tensor, ignore_index: int = 255, ): """ Métrica operacional: target/alvo = vegetação viva e não cana Usa heads vegetation e cana. Classe 0 = não alvo Classe 1 = alvo/pulverizável """ valid = (gt_veg != ignore_index) & (gt_cana != ignore_index) gt_target = ((gt_veg == 1) & (gt_cana == 0)).long() pred_target = ((pred_veg == 1) & (pred_cana == 0)).long() gt_target = gt_target[valid] pred_target = pred_target[valid] if gt_target.numel() == 0: return idx = gt_target.reshape(-1) * 2 + pred_target.reshape(-1) bins = torch.bincount(idx, minlength=4) cm_target += bins.view(2, 2) def compute_losses_for_batch( logits_by_head: Dict[str, torch.Tensor], masks_by_head: Dict[str, torch.Tensor], heads_config: Dict[str, dict], criterions: Dict[str, nn.Module], dice_weight: float = 0.30, config: Optional[dict] = None, epoch: Optional[int] = None, ): total = None loss_parts = {} cfg = config or {} target_distill_cfg = cfg.get("target_distillation", {}) or {} target_distill_enabled = bool(target_distill_cfg.get("enabled", False)) target_distill_ramp = get_target_distill_ramp_factor(cfg, epoch) target_distill_enabled = target_distill_enabled and target_distill_ramp > 0.0 for head_name, logits in logits_by_head.items(): target = masks_by_head[head_name] hcfg = heads_config[head_name] num_classes = int(hcfg["num_classes"]) ignore_index = int(hcfg.get("ignore_index", 255)) head_weight = float(hcfg.get("loss_weight_norm", 1.0)) if logits.shape[-2:] != target.shape[-2:]: logits = F.interpolate( logits, size=target.shape[-2:], mode="bilinear", align_corners=False, ) ce = criterions[head_name](logits, target) dl = dice_loss( logits=logits, target=target, num_classes=num_classes, ignore_index=ignore_index, smooth=1.0, ) head_loss = (1.0 - dice_weight) * ce + dice_weight * dl if head_name == "target" and target_distill_enabled: teacher = build_target_teacher_from_heads( logits_by_head=logits_by_head, target_shape=target.shape[-2:], config=cfg, erva_id=int( cfg.get("heads", {}) .get("semantic", {}) .get("classes", {}) .get("erva", 2) ), ) if teacher is not None: logits_target = logits if logits_target.shape[-2:] != target.shape[-2:]: logits_target = F.interpolate( logits_target, size=target.shape[-2:], mode="bilinear", align_corners=False, ) logits_binary = logits_target[:, 1, :, :] - logits_target[:, 0, :, :] valid = target != ignore_index if valid.any(): logits_binary_v = logits_binary[valid] teacher_v = teacher[valid].clamp(0.0, 1.0) distill = F.binary_cross_entropy_with_logits( logits_binary_v, teacher_v, ) hard_weight = float(target_distill_cfg.get("hard_weight", 0.70)) distill_weight = float(target_distill_cfg.get("distill_weight", 0.30)) # Aplica ramp-up só na parte destilada distill_weight = distill_weight * float(target_distill_ramp) denom = max(hard_weight + distill_weight, 1e-6) hard_weight = hard_weight / denom distill_weight = distill_weight / denom head_loss = hard_weight * head_loss + distill_weight * distill loss_parts[f"{head_name}_distill"] = distill.detach() loss_parts[f"{head_name}_distill_ramp"] = torch.as_tensor( target_distill_ramp, device=distill.device, dtype=distill.dtype, ).detach() weighted = head_weight * head_loss total = weighted if total is None else total + weighted loss_parts[f"{head_name}_loss"] = head_loss.detach() loss_parts[f"{head_name}_ce"] = ce.detach() loss_parts[f"{head_name}_dice"] = dl.detach() return total, loss_parts def resize_logits_to_target_if_needed( logits: torch.Tensor, target: torch.Tensor, ) -> torch.Tensor: """ Garante que logits estejam no mesmo HxW da máscara. Usado para métricas, sem depender do efeito colateral da loss. """ if logits.shape[-2:] != target.shape[-2:]: logits = F.interpolate( logits, size=target.shape[-2:], mode="bilinear", align_corners=False, ) return logits def get_target_distill_ramp_factor(config: dict, epoch: Optional[int]) -> float: cfg = (config or {}).get("target_distillation", {}) or {} if not bool(cfg.get("enabled", False)): return 0.0 if not bool(cfg.get("rampup_enabled", True)): return 1.0 if epoch is None: return 1.0 start_epoch = int(cfg.get("start_epoch", 8)) rampup_epochs = int(cfg.get("rampup_epochs", 12)) if epoch < start_epoch: return 0.0 if rampup_epochs <= 0: return 1.0 t = (float(epoch) - float(start_epoch) + 1.0) / float(rampup_epochs) t = max(0.0, min(1.0, t)) # Smoothstep: sobe suave, sem tranco return float(t * t * (3.0 - 2.0 * t)) # ============================================================ # Train / Val # ============================================================ def run_one_epoch( model, loader, optimizer, device, heads_config, criterions, amp, scaler, train, grad_accum=1, normalizer=None, dice_weight=0.30, config=None, epoch: Optional[int] = None, ): model.train(train) total_loss = 0.0 n_batches = 0 loss_sums = {} cms = { head_name: torch.zeros( (int(hcfg["num_classes"]), int(hcfg["num_classes"])), dtype=torch.int64, device=device, ) for head_name, hcfg in heads_config.items() } cm_target = torch.zeros((2, 2), dtype=torch.int64, device=device) t0 = time.time() with torch.set_grad_enabled(train): if train and optimizer is not None: optimizer.zero_grad(set_to_none=True) for step, (imgs, masks, _meta) in enumerate(loader): imgs = imgs.to(device, non_blocking=True) masks = {k: v.to(device, non_blocking=True) for k, v in masks.items()} if normalizer is not None: imgs = normalizer(imgs) else: imgs = normalize_per_batch(imgs) with autocast(device_type="cuda", enabled=amp and device.type == "cuda"): logits_by_head = model(pixel_values=imgs) loss, loss_parts = compute_losses_for_batch( logits_by_head=logits_by_head, masks_by_head=masks, heads_config=heads_config, criterions=criterions, dice_weight=dice_weight, config=config, epoch=epoch, ) 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 or (step + 1) == len(loader): scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) else: loss.backward() if ((step + 1) % grad_accum) == 0 or (step + 1) == len(loader): optimizer.step() optimizer.zero_grad(set_to_none=True) loss_value = float(loss.item()) * (grad_accum if train and grad_accum > 1 else 1.0) total_loss += loss_value n_batches += 1 for k, v in loss_parts.items(): loss_sums[k] = loss_sums.get(k, 0.0) + float(v.item()) with torch.no_grad(): preds_by_head = {} for head_name, logits in logits_by_head.items(): if head_name not in masks: continue target_mask = masks[head_name] logits_metric = resize_logits_to_target_if_needed( logits=logits, target=target_mask, ) preds = torch.argmax(logits_metric, dim=1) preds_by_head[head_name] = preds update_confusion_matrix( cms[head_name], preds, target_mask, int(heads_config[head_name]["num_classes"]), int(heads_config[head_name].get("ignore_index", 255)), ) # Métrica operacional antiga: # target_op = vegetation == 1 AND cana == 0 # Continua útil para comparar com a nova head target direta. if "vegetation" in preds_by_head and "cana" in preds_by_head: update_operational_target_cm( cm_target=cm_target, pred_veg=preds_by_head["vegetation"], pred_cana=preds_by_head["cana"], gt_veg=masks["vegetation"], gt_cana=masks["cana"], ignore_index=int(heads_config["vegetation"].get("ignore_index", 255)), ) avg_loss = total_loss / max(1, n_batches) metrics = { "loss": avg_loss, "time_s": time.time() - t0, "heads": {}, "loss_parts": {k: v / max(1, n_batches) for k, v in loss_sums.items()}, } for head_name, cm in cms.items(): miou, iou_per_class = compute_iou_from_cm(cm) acc = compute_pixel_acc_from_cm(cm) metrics["heads"][head_name] = { "miou": miou, "iou_per_class": iou_per_class, "acc": acc, "cm": cm.detach().cpu().tolist(), } target_miou, target_iou = compute_iou_from_cm(cm_target) target_acc = compute_pixel_acc_from_cm(cm_target) metrics["operational_target"] = { "miou": target_miou, "iou_background": target_iou[0], "iou_target": target_iou[1], "acc": target_acc, "cm": cm_target.detach().cpu().tolist(), } return metrics # ============================================================ # Checkpoint / logs # ============================================================ def save_checkpoint(path, model, optimizer, scaler, epoch, best: dict, extra=None): ckpt = { "epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict(), "best": best, } if scaler is not None: ckpt["scaler"] = scaler.state_dict() if extra: ckpt["extra"] = extra torch.save(ckpt, path) def load_checkpoint(path, model, optimizer=None, scaler=None, map_location="cpu"): 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 append_train_log(path: Path, row: dict): ensure_dir(path.parent) exists = path.exists() with path.open("a", newline="", encoding="utf-8") as f: w = csv.DictWriter(f, fieldnames=list(row.keys())) if not exists: w.writeheader() w.writerow(row) def flatten_epoch_log(epoch: int, lr: float, tr: dict, va: dict, best: dict) -> dict: row = { "epoch": epoch, "lr": lr, "train_loss": tr["loss"], "val_loss": va["loss"], "best_score": best.get("score", -1.0), "best_target_iou": best.get("target_iou", -1.0), "best_cana_head_iou": best.get("cana_head_iou", -1.0), "best_semantic_miou": best.get("semantic_miou", -1.0), "best_target_head_iou": best.get("target_head_iou", -1.0), "best_operational_target_iou": best.get("operational_target_iou", -1.0), } for prefix, obj in (("train", tr), ("val", va)): for head_name, hm in obj.get("heads", {}).items(): row[f"{prefix}_{head_name}_miou"] = hm.get("miou") row[f"{prefix}_{head_name}_acc"] = hm.get("acc") ious = hm.get("iou_per_class", []) or [] for i, v in enumerate(ious): row[f"{prefix}_{head_name}_iou_{i}"] = v op = obj.get("operational_target", {}) or {} row[f"{prefix}_target_miou"] = op.get("miou") row[f"{prefix}_target_iou"] = op.get("iou_target") row[f"{prefix}_target_acc"] = op.get("acc") for k, v in obj.get("loss_parts", {}).items(): row[f"{prefix}_{k}"] = v return row # ============================================================ # Pretty print # ============================================================ def pretty_iou(names: Dict[int, str], iou_list: List[float]): return " | ".join([ f"{names.get(i, i)}:{v:.3f}" for i, v in enumerate(iou_list) ]) def binary_iou_text(head_name: str, iou_list: List[float]): if head_name == "vegetation": names = {0: "bg", 1: "veg"} elif head_name == "cana": names = {0: "not_cana", 1: "cana"} elif head_name == "target": names = {0: "bg", 1: "target"} else: names = {0: "0", 1: "1"} return pretty_iou(names, iou_list) def compute_selection_score(va: dict) -> dict: heads = va.get("heads", {}) op = va.get("operational_target", {}) # Preferir target direta, se existir target_head_iou = 0.0 if "target" in heads: vals = heads["target"].get("iou_per_class", []) or [] if len(vals) > 1: target_head_iou = safe_float(vals[1], 0.0) operational_target_iou = safe_float(op.get("iou_target"), 0.0) # Se target head existir, ela manda. Se não existir, usa operacional antigo. target_iou = target_head_iou if "target" in heads else operational_target_iou cana_iou = 0.0 if "cana" in heads: vals = heads["cana"].get("iou_per_class", []) or [] if len(vals) > 1: cana_iou = safe_float(vals[1], 0.0) veg_miou = safe_float(heads.get("vegetation", {}).get("miou"), 0.0) semantic_miou = safe_float(heads.get("semantic", {}).get("miou"), 0.0) score = ( 0.50 * target_iou + 0.25 * cana_iou + 0.15 * veg_miou + 0.10 * semantic_miou ) return { "score": float(score), "target_iou": float(target_iou), "target_head_iou": float(target_head_iou), "operational_target_iou": float(operational_target_iou), "cana_head_iou": float(cana_iou), "vegetation_miou": float(veg_miou), "semantic_miou": float(semantic_miou), } # ============================================================ # Main # ============================================================ def main(): parser = argparse.ArgumentParser() parser.add_argument("--config", default="config.json") parser.add_argument("--epochs", type=int, default=80) parser.add_argument("--batch", type=int, default=1) parser.add_argument("--lr", type=float, default=3e-5) parser.add_argument("--wd", type=float, default=0.01) parser.add_argument("--num_workers", type=int, default=2) parser.add_argument("--amp", action="store_true") parser.add_argument("--amp_val", action="store_true") parser.add_argument("--grad_accum", type=int, default=4) parser.add_argument("--grad_ckpt", action="store_true") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--ignore_index", type=int, default=None) parser.add_argument("--class_weights", default="auto") parser.add_argument("--dice_weight", type=float, default=0.30) parser.add_argument("--norm_stats", default=None) parser.add_argument("--src-root", default="dataset/split") parser.add_argument("--save-every", type=int, default=10) parser.add_argument("--resume", action="store_true") parser.add_argument("--resume-ckpt", default=None) parser.add_argument("--early-stop", type=int, default=25) args = parser.parse_args() set_seed(args.seed) config = load_json(args.config) W, H = config["resolucao"] channels = int(config.get("channels", 5)) backbone = config.get("backbone", "nvidia/mit-b1") fusion_mode = config.get("fusion_mode", "stacked") input_channel_names = get_input_channel_names(config) input_channel_indices = get_input_channel_indices(config) channels = len(input_channel_names) print(f"Input channels: {input_channel_names} idx={input_channel_indices}") if fusion_mode != "stacked": raise RuntimeError("Este script é para fusion_mode='stacked'.") model_family = config.get("modelo", "segformer") model_name = config.get("model_name", "test") stats_source_tag = config.get("stats_source_tag", f"stacked_raw{channels}") save_dir = Path("backup") / model_family / model_name / f"{fusion_mode}_raw{channels}" ensure_dir(save_dir) labelmap_path = Path("dataset") / "labelmap.txt" if not labelmap_path.exists(): raise RuntimeError(f"Labelmap não encontrado: {labelmap_path}") semantic_id2label, semantic_label2id, ignore_from_labelmap = load_labelmap(str(labelmap_path)) ignore_index = int(args.ignore_index if args.ignore_index is not None else ignore_from_labelmap) heads_config = build_heads_config(config, ignore_index=ignore_index) # Garante que a semantic conhece o num_classes real do labelmap. heads_config["semantic"]["num_classes"] = int(len(semantic_id2label)) heads_config["semantic"]["ignore_index"] = int(ignore_index) print("==========================================") print("Train SegFormer OAK-FCC-3 Multi-Head") print(f"Backbone : {backbone}") print(f"Save dir : {save_dir}") print(f"Split root : {args.src_root}") print(f"Resolution : {W}x{H}") print(f"Channels : {channels}") print(f"Semantic : {heads_config['semantic']['num_classes']} -> {semantic_id2label}") print(f"Ignore index : {ignore_index}") print("Heads:") for name, hcfg in heads_config.items(): print( f" - {name}: classes={hcfg['num_classes']} " f"mask_dir={hcfg['mask_dir']} " f"loss_weight={hcfg['loss_weight']} norm={hcfg['loss_weight_norm']:.3f}" ) print("==========================================") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") train_root = Path(args.src_root) / "train" val_root = Path(args.src_root) / "val" resize_hw = (H, W) ds_train = OakFcc3TensorMultiHeadDataset( train_root, heads_config=heads_config, channels=channels, channel_indices=input_channel_indices, strict_channels=False, resize_hw=resize_hw, ) ds_val = OakFcc3TensorMultiHeadDataset( val_root, heads_config=heads_config, channels=channels, channel_indices=input_channel_indices, strict_channels=False, resize_hw=resize_hw, ) print(f"[DATA] train={len(ds_train)} | val={len(ds_val)}") dl_train = DataLoader( ds_train, batch_size=args.batch, shuffle=True, num_workers=args.num_workers, pin_memory=True, collate_fn=collate_fn, drop_last=True if len(ds_train) >= args.batch else False, ) dl_val = DataLoader( ds_val, batch_size=1, shuffle=False, num_workers=max(0, args.num_workers // 2), pin_memory=True, collate_fn=collate_fn, drop_last=False, ) normalizer, norm_stats_path = build_normalizer(config, args, device) model = build_model( backbone=backbone, channels=channels, heads_config=heads_config, semantic_id2label=semantic_id2label, semantic_label2id=semantic_label2id, ) if args.grad_ckpt: try: model.segformer.gradient_checkpointing_enable() print("[MODEL] gradient checkpointing enabled") except Exception as e: print(f"[WARN] gradient checkpointing não suportado: {e}") model.to(device) criterions, weight_debug = build_criterions( ds_train=ds_train, heads_config=heads_config, class_weights_mode=args.class_weights, device=device, seed=args.seed, ) 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, ) scaler = GradScaler(enabled=args.amp and device.type == "cuda") last_path = save_dir / "last.pt" best_score_path = save_dir / "best_score.pt" best_target_path = save_dir / "best_target.pt" best_cana_path = save_dir / "best_cana_head.pt" best_semantic_path = save_dir / "best_semantic_miou.pt" train_log_path = save_dir / "train_log.csv" save_json(save_dir / "train_config_snapshot.json", { "config": config, "args": vars(args), "semantic_id2label": semantic_id2label, "semantic_label2id": semantic_label2id, "heads_config": heads_config, "ignore_index": ignore_index, "norm_stats_path": norm_stats_path, "class_weight_debug": weight_debug, "input_channel_names": input_channel_names, "input_channel_indices": input_channel_indices, "checkpoint_selection_score": { "score": "0.45*target_iou + 0.30*cana_head_iou + 0.15*vegetation_miou + 0.10*semantic_miou" }, }) start_epoch = 1 best = { "score": -1.0, "target_iou": -1.0, "target_head_iou": -1.0, "operational_target_iou": -1.0, "cana_head_iou": -1.0, "semantic_miou": -1.0, "vegetation_miou": -1.0, } epochs_without_improve = 0 resume_path = Path(args.resume_ckpt) if args.resume_ckpt else last_path if args.resume and resume_path.exists(): ckpt = load_checkpoint( resume_path, model, optimizer, scaler=scaler, map_location="cpu", ) start_epoch = int(ckpt["epoch"]) + 1 best = ckpt.get("best", best) print(f"[RESUME] {resume_path} epoch={start_epoch}") 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, heads_config=heads_config, criterions=criterions, amp=args.amp, scaler=scaler, train=True, grad_accum=max(1, args.grad_accum), normalizer=normalizer, dice_weight=args.dice_weight, config=config, epoch=epoch, ) if device.type == "cuda": torch.cuda.empty_cache() va = run_one_epoch( model=model, loader=dl_val, optimizer=None, device=device, heads_config=heads_config, criterions=criterions, amp=args.amp_val, scaler=None, train=False, grad_accum=1, normalizer=normalizer, dice_weight=args.dice_weight, config=config, epoch=epoch, ) scheduler.step(va["loss"]) score_now = compute_selection_score(va) print( f"TRAIN: loss={tr['loss']:.4f} " f"sem_mIoU={tr['heads']['semantic']['miou']:.4f} " f"veg_mIoU={tr['heads']['vegetation']['miou']:.4f} " f"cana_mIoU={tr['heads']['cana']['miou']:.4f} " f"targetIoU={tr['operational_target']['iou_target']:.4f} " f"t={tr['time_s']:.1f}s" ) target_head_iou_val = 0.0 if "target" in va["heads"]: vals = va["heads"]["target"].get("iou_per_class", []) or [] if len(vals) > 1: target_head_iou_val = vals[1] print( f"VAL : loss={va['loss']:.4f} " f"sem_mIoU={va['heads']['semantic']['miou']:.4f} " f"veg_mIoU={va['heads']['vegetation']['miou']:.4f} " f"cana_mIoU={va['heads']['cana']['miou']:.4f} " f"targetHeadIoU={target_head_iou_val:.4f} " f"opTargetIoU={va['operational_target']['iou_target']:.4f} " f"score={score_now['score']:.4f} " f"t={va['time_s']:.1f}s" ) print("IoU semantic:", pretty_iou(semantic_id2label, va["heads"]["semantic"]["iou_per_class"])) print("IoU vegetation:", binary_iou_text("vegetation", va["heads"]["vegetation"]["iou_per_class"])) print("IoU cana_head:", binary_iou_text("cana", va["heads"]["cana"]["iou_per_class"])) print("IoU target:", f"bg:{va['operational_target']['iou_background']:.3f} | alvo:{va['operational_target']['iou_target']:.3f}") if "target" in va["heads"]: print("IoU target_head:", binary_iou_text("target", va["heads"]["target"]["iou_per_class"])) save_checkpoint( last_path, model, optimizer, scaler=scaler, epoch=epoch, best=best, extra={ "train": tr, "val": va, "score_now": score_now, }, ) if args.save_every > 0 and epoch % args.save_every == 0: save_checkpoint( save_dir / f"epoch_{epoch:04d}.pt", model, optimizer, scaler=scaler, epoch=epoch, best=best, ) improved = False if score_now["score"] > best.get("score", -1.0): best["score"] = score_now["score"] improved = True save_checkpoint( best_score_path, model, optimizer, scaler=scaler, epoch=epoch, best=best, extra={"val": va, "score_now": score_now}, ) print(f"[BEST SCORE] {best['score']:.4f} -> {best_score_path}") if score_now["target_iou"] > best.get("target_iou", -1.0): best["target_iou"] = score_now["target_iou"] improved = True save_checkpoint( best_target_path, model, optimizer, scaler=scaler, epoch=epoch, best=best, extra={"val": va, "score_now": score_now}, ) print(f"[BEST TARGET] {best['target_iou']:.4f} -> {best_target_path}") if score_now["cana_head_iou"] > best.get("cana_head_iou", -1.0): best["cana_head_iou"] = score_now["cana_head_iou"] improved = True save_checkpoint( best_cana_path, model, optimizer, scaler=scaler, epoch=epoch, best=best, extra={"val": va, "score_now": score_now}, ) print(f"[BEST CANA HEAD] {best['cana_head_iou']:.4f} -> {best_cana_path}") if score_now["semantic_miou"] > best.get("semantic_miou", -1.0): best["semantic_miou"] = score_now["semantic_miou"] improved = True save_checkpoint( best_semantic_path, model, optimizer, scaler=scaler, epoch=epoch, best=best, extra={"val": va, "score_now": score_now}, ) print(f"[BEST SEMANTIC] {best['semantic_miou']:.4f} -> {best_semantic_path}") append_train_log( train_log_path, flatten_epoch_log( epoch=epoch, lr=lr_now, tr=tr, va=va, best=best, ), ) if improved: epochs_without_improve = 0 else: epochs_without_improve += 1 if args.early_stop > 0 and epochs_without_improve >= args.early_stop: print(f"[EARLY STOP] {epochs_without_improve} épocas sem melhora.") break print("\nTreino finalizado.") print(f"Best score : {best.get('score', -1.0):.4f}") print(f"Best target IoU : {best.get('target_iou', -1.0):.4f}") print(f"Best cana-head IoU: {best.get('cana_head_iou', -1.0):.4f}") print(f"Best semantic mIoU: {best.get('semantic_miou', -1.0):.4f}") print(f"Save dir : {save_dir}") if __name__ == "__main__": main()