700 lines
28 KiB
Python
700 lines
28 KiB
Python
|
# 👉 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)
|