import os import json import time import copy from pathlib import Path from typing import Optional, Dict import numpy as np import torch import torch.nn as nn from torch.utils.data import DataLoader from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts DEFAULT_CONFIG = { "arch": "hrnet", "num_outputs": 3, "image_size": 512, "batch_size": 8, "num_workers": 4, "epochs": 100, "lr": 1e-3, "weight_decay": 1e-4, "patience": 15, "t0": 10, "t_mult": 2, "eta_min": 1e-6, "output_dir": "results/", "exp_name": "scoliosis_hrnet", } def train_one_epoch(model, loader, criterion, optimizer, device): model.train() losses, maes = [], [] for batch in loader: images = batch["image"].to(device, non_blocking=True) targets = batch["angles"].to(device, non_blocking=True) preds = model(images) loss = criterion(preds, targets) optimizer.zero_grad(set_to_none=True) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0) optimizer.step() losses.append(loss.item()) maes.append(torch.abs(preds.detach() - targets).mean().item()) return {"loss": np.mean(losses), "cmae": np.mean(maes)} @torch.no_grad() def validate(model, loader, criterion, device): model.eval() losses, maes = [], [] all_pred, all_target = [], [] for batch in loader: images = batch["image"].to(device, non_blocking=True) targets = batch["angles"].to(device, non_blocking=True) preds = model(images) loss = criterion(preds, targets) losses.append(loss.item()) maes.append(torch.abs(preds - targets).mean().item()) all_pred.append(preds.cpu()) all_target.append(targets.cpu()) all_pred = torch.cat(all_pred) all_target = torch.cat(all_target) per_angle = torch.abs(all_pred - all_target).mean(dim=0) return { "loss": np.mean(losses), "cmae": np.mean(maes), "cmae_thoracic1": per_angle[0].item(), "cmae_thoracic2": per_angle[1].item(), "cmae_lumbar": per_angle[2].item(), } def train(model, train_loader, val_loader, cfg, resume_ckpt=None): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"\n[Treino] Dispositivo: {device}") model = model.to(device) out_dir = Path(cfg["output_dir"]) / cfg["exp_name"] out_dir.mkdir(parents=True, exist_ok=True) with open(out_dir / "config.json", "w") as f: json.dump(cfg, f, indent=2) optimizer = AdamW(model.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"]) scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=cfg["t0"], T_mult=cfg["t_mult"], eta_min=cfg["eta_min"]) criterion = nn.SmoothL1Loss() best_cmae = float("inf") best_epoch = 0 best_weights = None patience_ctr = 0 history = {"train": [], "val": []} start_epoch = 0 if resume_ckpt and os.path.isfile(resume_ckpt): ckpt = torch.load(resume_ckpt, map_location=device) model.load_state_dict(ckpt["model"]) optimizer.load_state_dict(ckpt["optimizer"]) scheduler.load_state_dict(ckpt["scheduler"]) start_epoch = ckpt["epoch"] + 1 best_cmae = ckpt.get("best_cmae", float("inf")) history = ckpt.get("history", history) print(f"[Treino] Retomando do epoch {start_epoch} | best CMAE: {best_cmae:.2f}°") for epoch in range(start_epoch, cfg["epochs"]): t0 = time.time() print(f"\nEpoch [{epoch+1}/{cfg['epochs']}] — lr={scheduler.get_last_lr()[0]:.2e}") train_m = train_one_epoch(model, train_loader, criterion, optimizer, device) val_m = validate(model, val_loader, criterion, device) scheduler.step() history["train"].append(train_m) history["val"].append(val_m) print(f" TREINO — loss={train_m['loss']:.4f} | CMAE={train_m['cmae']:.2f}deg") print(f" VAL — loss={val_m['loss']:.4f} | CMAE={val_m['cmae']:.2f}deg " f"[T1={val_m['cmae_thoracic1']:.1f} T2={val_m['cmae_thoracic2']:.1f} L={val_m['cmae_lumbar']:.1f}] | {time.time()-t0:.1f}s") with open(out_dir / "history.json", "w") as f: json.dump(history, f, indent=2) torch.save({ "epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), "best_cmae": best_cmae, "history": history, "cfg": cfg, }, out_dir / "last.pth") if val_m["cmae"] < best_cmae: best_cmae = val_m["cmae"] best_epoch = epoch + 1 best_weights = copy.deepcopy(model.state_dict()) patience_ctr = 0 torch.save({"epoch": epoch, "model": best_weights, "cmae": best_cmae, "cfg": cfg}, out_dir / "best.pth") print(f" ✓ Novo melhor modelo — CMAE={best_cmae:.2f}° (epoch {best_epoch})") else: patience_ctr += 1 if patience_ctr >= cfg["patience"]: print(f"\n[Early stopping] Sem melhoria há {cfg['patience']} epochs.") break print(f"\n[Treino concluído] Melhor CMAE: {best_cmae:.2f}° (epoch {best_epoch})") if best_weights: model.load_state_dict(best_weights) return history @torch.no_grad() def evaluate_test(model, test_loader, cfg, ckpt_path=None): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) if ckpt_path: ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) model.load_state_dict(ckpt["model"]) print(f"[Teste] Checkpoint: {ckpt_path}") criterion = nn.SmoothL1Loss() results = validate(model, test_loader, criterion, device) print("\n── Resultados no conjunto de teste ──") print(f" CMAE medio: {results['cmae']:.2f} graus") print(f" CMAE toracico (T1): {results['cmae_thoracic1']:.2f} graus") print(f" CMAE toracico (T2): {results['cmae_thoracic2']:.2f} graus") print(f" CMAE lombar: {results['cmae_lumbar']:.2f} graus") return results