fix: corrige duplicação de diretório do treino e valida ambiente contra allowlist
O fix anterior (forçar project="runs/detect") duplicava o caminho para
"runs/detect/runs/detect/{ambiente}" porque o Ultralytics já resolve esse
project por padrão a partir do runs_dir configurado. Resultado: o "best"
nunca era encontrado e NENHUM modelo era enviado ao S3 (pior que o bug
anterior, que ao menos subia um modelo antigo). Agora só passamos
name/exist_ok e lemos o caminho real via modelo_base.trainer.save_dir,
sem adivinhar a string.
Também adiciona allowlist para 'ambiente' em /detectar e /treinar — o
valor é usado sem sanitização para montar paths locais (/tmp/...,
runs/detect/...) e chaves S3, o que permitia path traversal.
This commit is contained in:
@@ -14,6 +14,11 @@ app = FastAPI()
|
|||||||
s3 = boto3.client('s3')
|
s3 = boto3.client('s3')
|
||||||
BUCKET = os.getenv('BUCKET_NAME', 'ia-gondola-projeto-2024')
|
BUCKET = os.getenv('BUCKET_NAME', 'ia-gondola-projeto-2024')
|
||||||
|
|
||||||
|
# 'ambiente' vira componente de paths locais (/tmp/..., runs/detect/...) e de
|
||||||
|
# chaves S3 — sem allowlist, um valor tipo "../../../app" permitiria path
|
||||||
|
# traversal (leitura/escrita/exclusão fora do diretório esperado).
|
||||||
|
AMBIENTES_VALIDOS = {"gondola", "etiqueta"}
|
||||||
|
|
||||||
MODELOS_CARREGADOS = {}
|
MODELOS_CARREGADOS = {}
|
||||||
|
|
||||||
_treinamento_status = {"status": "idle", "ambiente": None, "versao": None, "detalhe": None}
|
_treinamento_status = {"status": "idle", "ambiente": None, "versao": None, "detalhe": None}
|
||||||
@@ -187,18 +192,19 @@ def _executar_treino_sync(ambiente: str, pular_triagem: bool, forcar_do_zero: bo
|
|||||||
modelo_base.add_callback("on_fit_epoch_end", _on_fit_epoch_end)
|
modelo_base.add_callback("on_fit_epoch_end", _on_fit_epoch_end)
|
||||||
# Fine-tuning: LR baixo preserva o que o modelo já aprendeu.
|
# Fine-tuning: LR baixo preserva o que o modelo já aprendeu.
|
||||||
# Primeiro treino: LR padrão + mais épocas para convergir do zero.
|
# Primeiro treino: LR padrão + mais épocas para convergir do zero.
|
||||||
# project/name/exist_ok fixos: sem isso o Ultralytics incrementa o
|
# name/exist_ok fixos: sem isso o Ultralytics incrementa o diretório a
|
||||||
# diretório a cada treino (train, train2, train3...) e o caminho abaixo
|
# cada treino (train, train2, train3...). Não passar "project" — o
|
||||||
# ficava sempre apontando pro "train" original, subindo pro S3 um
|
# Ultralytics já resolve o project default a partir do runs_dir das
|
||||||
# modelo de uma rodada antiga em vez do que acabou de treinar.
|
# settings, e forçar um valor aqui duplicava o caminho
|
||||||
train_dir = f"runs/detect/{ambiente}"
|
# (runs/detect/runs/detect/{ambiente}), fazendo o `best` abaixo nunca
|
||||||
|
# ser encontrado. Lemos o caminho real direto do trainer, sem adivinhar.
|
||||||
if is_fine_tuning:
|
if is_fine_tuning:
|
||||||
modelo_base.train(data=yaml_path, epochs=20, imgsz=640, batch=8, device='cpu', plots=True, lr0=0.001, lrf=0.01, project="runs/detect", name=ambiente, exist_ok=True)
|
modelo_base.train(data=yaml_path, epochs=20, imgsz=640, batch=8, device='cpu', plots=True, lr0=0.001, lrf=0.01, name=ambiente, exist_ok=True)
|
||||||
else:
|
else:
|
||||||
modelo_base.train(data=yaml_path, epochs=30, imgsz=640, batch=8, device='cpu', plots=True, lr0=0.01, lrf=0.01, project="runs/detect", name=ambiente, exist_ok=True)
|
modelo_base.train(data=yaml_path, epochs=30, imgsz=640, batch=8, device='cpu', plots=True, lr0=0.01, lrf=0.01, name=ambiente, exist_ok=True)
|
||||||
modelo_base.reset_callbacks()
|
modelo_base.reset_callbacks()
|
||||||
|
|
||||||
best = f"{train_dir}/weights/best.pt"
|
best = str(modelo_base.trainer.save_dir / "weights" / "best.pt")
|
||||||
if os.path.exists(best):
|
if os.path.exists(best):
|
||||||
carimbo = datetime.now().strftime("%Y%m%d_%H%M")
|
carimbo = datetime.now().strftime("%Y%m%d_%H%M")
|
||||||
s3.upload_file(best, BUCKET, f"modelos/{ambiente}/atual/cerebro.pt")
|
s3.upload_file(best, BUCKET, f"modelos/{ambiente}/atual/cerebro.pt")
|
||||||
@@ -246,6 +252,8 @@ def _treinar_bg(ambiente: str, pular_triagem: bool, forcar_do_zero: bool = False
|
|||||||
|
|
||||||
@app.post("/detectar")
|
@app.post("/detectar")
|
||||||
async def detectar(ambiente: str = Form(...), file: UploadFile = File(...)):
|
async def detectar(ambiente: str = Form(...), file: UploadFile = File(...)):
|
||||||
|
if ambiente not in AMBIENTES_VALIDOS:
|
||||||
|
raise HTTPException(status_code=400, detail=f"ambiente inválido: {ambiente}")
|
||||||
try:
|
try:
|
||||||
Image.MAX_IMAGE_PIXELS = None
|
Image.MAX_IMAGE_PIXELS = None
|
||||||
modelo = carregar_modelo_do_s3(ambiente)
|
modelo = carregar_modelo_do_s3(ambiente)
|
||||||
@@ -265,10 +273,12 @@ async def detectar(ambiente: str = Form(...), file: UploadFile = File(...)):
|
|||||||
@app.post("/treinar")
|
@app.post("/treinar")
|
||||||
async def treinar(dados: dict):
|
async def treinar(dados: dict):
|
||||||
global _treinamento_status
|
global _treinamento_status
|
||||||
|
ambiente = dados.get("ambiente", "gondola")
|
||||||
|
if ambiente not in AMBIENTES_VALIDOS:
|
||||||
|
raise HTTPException(status_code=400, detail=f"ambiente inválido: {ambiente}")
|
||||||
with _status_lock:
|
with _status_lock:
|
||||||
if _treinamento_status["status"] == "running":
|
if _treinamento_status["status"] == "running":
|
||||||
raise HTTPException(status_code=409, detail="Treinamento já em andamento")
|
raise HTTPException(status_code=409, detail="Treinamento já em andamento")
|
||||||
ambiente = dados.get("ambiente", "gondola")
|
|
||||||
pular_triagem = dados.get("pular_triagem", False)
|
pular_triagem = dados.get("pular_triagem", False)
|
||||||
forcar_do_zero = dados.get("forcar_do_zero", False)
|
forcar_do_zero = dados.get("forcar_do_zero", False)
|
||||||
_treinamento_status = {"status": "running", "ambiente": ambiente, "versao": None, "detalhe": None}
|
_treinamento_status = {"status": "running", "ambiente": ambiente, "versao": None, "detalhe": None}
|
||||||
|
|||||||
Reference in New Issue
Block a user