diff --git a/main.py b/main.py index a62766c..1f4d0da 100644 --- a/main.py +++ b/main.py @@ -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__":