diff --git a/main.py b/main.py index 1f4d0da..cf82cee 100644 --- a/main.py +++ b/main.py @@ -83,7 +83,15 @@ def _executar_treino_sync(ambiente: str, pular_triagem: bool) -> dict: prefix_novos = f"treinamento/{ambiente}/novos-treinamentos/" objs = s3.list_objects_v2(Bucket=BUCKET, Prefix=prefix_novos) - modelo_base = carregar_modelo_do_s3(ambiente) + # Tenta carregar modelo atual do S3; se não existir, usa YOLO genérico (primeiro treino) + try: + modelo_base = carregar_modelo_do_s3(ambiente) + is_fine_tuning = True + except Exception: + log_print(f"Nenhum modelo base encontrado para '{ambiente}', iniciando do zero com yolov8n.") + modelo_base = YOLO('yolov8n.pt') + is_fine_tuning = False + pular_triagem = True # triagem sem modelo próprio é inútil if pular_triagem: log_print(f"TRIAGEM PULADA: movendo todos os arquivos direto para base-oficial ({ambiente})") @@ -143,9 +151,14 @@ def _executar_treino_sync(ambiente: str, pular_triagem: bool) -> dict: with open(yaml_path, 'w') as f: 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...") + log_print(f"Treinando com {len(os.listdir(img_dir))} fotos ({'fine-tuning' if is_fine_tuning else 'do zero'})...") 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) + # Fine-tuning: LR baixo preserva o que o modelo já aprendeu. + # Primeiro treino: LR padrão + mais épocas para convergir do zero. + if is_fine_tuning: + modelo_base.train(data=yaml_path, epochs=20, imgsz=640, batch=16, device='cpu', plots=True, lr0=0.001, lrf=0.01) + else: + modelo_base.train(data=yaml_path, epochs=30, imgsz=640, batch=16, device='cpu', plots=True, lr0=0.01, lrf=0.01) modelo_base.reset_callbacks() best = "runs/detect/train/weights/best.pt"