feat: captura de logs por época via callback YOLO e expõe em GET /treinar/status

This commit is contained in:
2026-06-10 15:14:10 -03:00
parent f7ab4219ca
commit 336b89cc98
+25 -1
View File
@@ -18,6 +18,26 @@ MODELOS_CARREGADOS = {}
_treinamento_status = {"status": "idle", "ambiente": None, "versao": None, "detalhe": None}
_status_lock = threading.Lock()
_training_logs: list = []
def _on_fit_epoch_end(trainer):
epoch = trainer.epoch + 1
total = trainer.epochs
try:
losses = [float(x) for x in trainer.loss_items] if getattr(trainer, 'loss_items', None) is not None else []
m = getattr(trainer, 'metrics', None) or {}
map50 = float(m.get('metrics/mAP50(B)', 0))
if len(losses) >= 3:
line = f"[{epoch}/{total}] box={losses[0]:.3f} cls={losses[1]:.3f} dfl={losses[2]:.3f} | mAP50={map50:.4f}"
else:
line = f"[{epoch}/{total}] mAP50={map50:.4f}"
except Exception:
line = f"[{epoch}/{total}]"
with _status_lock:
_training_logs.append(line)
if len(_training_logs) > 60:
_training_logs.pop(0)
def redimensionar_imagem(caminho):
@@ -124,7 +144,9 @@ def _executar_treino_sync(ambiente: str, pular_triagem: bool) -> dict:
yaml.dump({'train': img_dir, 'val': img_dir, 'nc': 1, 'names': {0: ambiente}}, f)
log_print(f"Treinando com {len(os.listdir(img_dir))} fotos...")
modelo_base.add_callback("on_fit_epoch_end", _on_fit_epoch_end)
modelo_base.train(data=yaml_path, epochs=30, imgsz=640, batch=16, device='cpu', plots=True)
modelo_base.reset_callbacks()
best = "runs/detect/train/weights/best.pt"
if os.path.exists(best):
@@ -142,6 +164,8 @@ def _executar_treino_sync(ambiente: str, pular_triagem: bool) -> dict:
def _treinar_bg(ambiente: str, pular_triagem: bool):
global _treinamento_status
with _status_lock:
_training_logs.clear()
try:
resultado = _executar_treino_sync(ambiente, pular_triagem)
with _status_lock:
@@ -205,7 +229,7 @@ async def treinar(dados: dict):
@app.get("/treinar/status")
async def status_treino():
with _status_lock:
return dict(_treinamento_status)
return {**dict(_treinamento_status), "logs": list(_training_logs[-30:])}
if __name__ == "__main__":