299 lines
13 KiB
Python
299 lines
13 KiB
Python
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
|
|
from ultralytics import YOLO
|
|
import boto3
|
|
import os
|
|
import io
|
|
import shutil
|
|
import yaml
|
|
import threading
|
|
from PIL import Image
|
|
from PIL import Image as PILImage
|
|
from datetime import datetime
|
|
|
|
app = FastAPI()
|
|
s3 = boto3.client('s3')
|
|
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 = {}
|
|
|
|
_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):
|
|
PILImage.MAX_IMAGE_PIXELS = None
|
|
img = PILImage.open(caminho)
|
|
w, h = img.size
|
|
max_px = 4096
|
|
if w > max_px or h > max_px:
|
|
ratio = min(max_px/w, max_px/h)
|
|
novo_w = int(w * ratio)
|
|
novo_h = int(h * ratio)
|
|
img = img.resize((novo_w, novo_h), PILImage.LANCZOS)
|
|
img.save(caminho, quality=95)
|
|
|
|
def log_print(msg):
|
|
carimbo = datetime.now().strftime("%H:%M:%S")
|
|
line = f"[{carimbo}] {msg}"
|
|
print(line, flush=True)
|
|
with _status_lock:
|
|
_training_logs.append(line)
|
|
if len(_training_logs) > 60:
|
|
_training_logs.pop(0)
|
|
|
|
def carregar_modelo_do_s3(ambiente: str):
|
|
s3_key = f"modelos/{ambiente}/atual/cerebro.pt"
|
|
local_path = f"/tmp/cerebro_{ambiente}.pt"
|
|
|
|
if os.path.exists(local_path):
|
|
data_local = os.path.getmtime(local_path)
|
|
resp_head = s3.head_object(Bucket=BUCKET, Key=s3_key)
|
|
data_s3 = resp_head['LastModified'].timestamp()
|
|
if data_s3 > data_local:
|
|
log_print(f"Modelo '{ambiente}' atualizado no S3. Atualizando local...")
|
|
os.remove(local_path)
|
|
if ambiente in MODELOS_CARREGADOS:
|
|
del MODELOS_CARREGADOS[ambiente]
|
|
|
|
if ambiente not in MODELOS_CARREGADOS:
|
|
if not os.path.exists(local_path):
|
|
log_print(f"Baixando cerebro '{ambiente}'...")
|
|
s3.download_file(BUCKET, s3_key, local_path)
|
|
MODELOS_CARREGADOS[ambiente] = YOLO(local_path)
|
|
return MODELOS_CARREGADOS[ambiente]
|
|
|
|
|
|
def _executar_treino_sync(ambiente: str, pular_triagem: bool, forcar_do_zero: bool = False) -> dict:
|
|
def _executar_treino_sync(ambiente: str, pular_triagem: bool, forcar_do_zero: bool = False) -> dict:
|
|
"""Execução síncrona do treino — chamada em background thread."""
|
|
prefix_novos = f"treinamento/{ambiente}/novos-treinamentos/"
|
|
objs = s3.list_objects_v2(Bucket=BUCKET, Prefix=prefix_novos)
|
|
|
|
if forcar_do_zero:
|
|
# Ignora o modelo atual (ex.: está degradado) sem tocar no arquivo em
|
|
# produção — /detectar continua servindo do cerebro.pt existente.
|
|
log_print(f"Treino do zero forçado para '{ambiente}' — ignorando modelo atual.")
|
|
modelo_base = YOLO('yolov8n.pt')
|
|
is_fine_tuning = False
|
|
pular_triagem = True # sem modelo próprio confiável, triagem não faz sentido
|
|
else:
|
|
# 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})")
|
|
if 'Contents' in objs:
|
|
for obj in objs['Contents']:
|
|
k = obj['Key']
|
|
if k.endswith(('.jpg', '.jpeg', '.png', '.txt')):
|
|
new_k = k.replace("novos-treinamentos", "base-oficial")
|
|
s3.copy_object(Bucket=BUCKET, CopySource={'Bucket': BUCKET, 'Key': k}, Key=new_k)
|
|
s3.delete_object(Bucket=BUCKET, Key=k)
|
|
log_print(f"{os.path.basename(k)} -> base-oficial")
|
|
else:
|
|
log_print(f"INICIANDO TRIAGEM INTELIGENTE: {ambiente}")
|
|
if 'Contents' in objs:
|
|
for obj in objs['Contents']:
|
|
if obj['Key'].endswith(('.jpg', '.jpeg', '.png')):
|
|
img_key = obj['Key']
|
|
txt_key = img_key.rsplit('.', 1)[0] + ".txt"
|
|
resp_img = s3.get_object(Bucket=BUCKET, Key=img_key)
|
|
img_data = Image.open(io.BytesIO(resp_img['Body'].read()))
|
|
pred = modelo_base(img_data, conf=0.1, verbose=False)
|
|
confs = [float(b.conf) for r in pred for b in r.boxes]
|
|
media = sum(confs)/len(confs) if confs else 0
|
|
decisao = "base-oficial" if media >= 0.15 else "descartados"
|
|
log_print(f"{os.path.basename(img_key)}: Conf. {media:.2f} -> {decisao}")
|
|
for k in [img_key, txt_key]:
|
|
try:
|
|
s3.head_object(Bucket=BUCKET, Key=k)
|
|
new_k = k.replace("novos-treinamentos", decisao)
|
|
s3.copy_object(Bucket=BUCKET, CopySource={'Bucket': BUCKET, 'Key': k}, Key=new_k)
|
|
s3.delete_object(Bucket=BUCKET, Key=k)
|
|
except Exception:
|
|
continue
|
|
|
|
dataset_local = f"/tmp/dataset_{ambiente}"
|
|
img_dir = f"{dataset_local}/train/images"
|
|
lbl_dir = f"{dataset_local}/train/labels"
|
|
if os.path.exists(dataset_local):
|
|
shutil.rmtree(dataset_local)
|
|
os.makedirs(img_dir, exist_ok=True)
|
|
os.makedirs(lbl_dir, exist_ok=True)
|
|
|
|
log_print("Baixando Base Oficial (Ouro)...")
|
|
ouro = s3.list_objects_v2(Bucket=BUCKET, Prefix=f"treinamento/{ambiente}/base-oficial/")
|
|
if 'Contents' in ouro:
|
|
for o in ouro['Contents']:
|
|
k = o['Key']
|
|
if k.endswith(('.jpg', '.jpeg', '.txt')):
|
|
dest = img_dir if not k.endswith('.txt') else lbl_dir
|
|
caminho_local = os.path.join(dest, os.path.basename(k))
|
|
s3.download_file(BUCKET, k, caminho_local)
|
|
if not k.endswith('.txt'):
|
|
redimensionar_imagem(caminho_local)
|
|
|
|
imagens = [f for f in os.listdir(img_dir) if not f.endswith('.txt')]
|
|
if len(imagens) > 0:
|
|
import random, shutil
|
|
random.shuffle(imagens)
|
|
corte = max(1, int(len(imagens) * 0.8))
|
|
val_imgs = imagens[corte:]
|
|
|
|
val_img_dir = f"{dataset_local}/val/images"
|
|
val_lbl_dir = f"{dataset_local}/val/labels"
|
|
os.makedirs(val_img_dir, exist_ok=True)
|
|
os.makedirs(val_lbl_dir, exist_ok=True)
|
|
|
|
for fname in val_imgs:
|
|
shutil.move(os.path.join(img_dir, fname), os.path.join(val_img_dir, fname))
|
|
lbl = fname.rsplit('.', 1)[0] + '.txt'
|
|
lbl_src = os.path.join(lbl_dir, lbl)
|
|
if os.path.exists(lbl_src):
|
|
shutil.move(lbl_src, os.path.join(val_lbl_dir, lbl))
|
|
|
|
yaml_path = f"{dataset_local}/data.yaml"
|
|
with open(yaml_path, 'w') as f:
|
|
yaml.dump({'train': img_dir, 'val': val_img_dir, 'nc': 1, 'names': {0: ambiente}}, f)
|
|
|
|
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)
|
|
# Fine-tuning: LR baixo preserva o que o modelo já aprendeu.
|
|
# Primeiro treino: LR padrão + mais épocas para convergir do zero.
|
|
# name/exist_ok fixos: sem isso o Ultralytics incrementa o diretório a
|
|
# cada treino (train, train2, train3...). Não passar "project" — o
|
|
# Ultralytics já resolve o project default a partir do runs_dir das
|
|
# settings, e forçar um valor aqui duplicava o caminho
|
|
# (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:
|
|
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:
|
|
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()
|
|
|
|
best = str(modelo_base.trainer.save_dir / "weights" / "best.pt")
|
|
if os.path.exists(best):
|
|
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}/versionamento/cerebro_{carimbo}.pt")
|
|
if ambiente in MODELOS_CARREGADOS:
|
|
del MODELOS_CARREGADOS[ambiente]
|
|
log_print("TREINAMENTO CONCLUIDO!")
|
|
return {"status": "sucesso", "versao": carimbo}
|
|
|
|
log_print("Nenhuma imagem passou na triagem para a Base Oficial.")
|
|
return {"status": "vazio"}
|
|
|
|
|
|
def _treinar_bg(ambiente: str, pular_triagem: bool, forcar_do_zero: bool = False):
|
|
global _treinamento_status
|
|
with _status_lock:
|
|
_training_logs.clear()
|
|
try:
|
|
resultado = _executar_treino_sync(ambiente, pular_triagem, forcar_do_zero)
|
|
with _status_lock:
|
|
if resultado["status"] == "sucesso":
|
|
_treinamento_status = {
|
|
"status": "concluido",
|
|
"ambiente": ambiente,
|
|
"versao": resultado.get("versao"),
|
|
"detalhe": None,
|
|
}
|
|
else:
|
|
_treinamento_status = {
|
|
"status": "vazio",
|
|
"ambiente": ambiente,
|
|
"versao": None,
|
|
"detalhe": "Nenhuma imagem passou na triagem",
|
|
}
|
|
except Exception as e:
|
|
log_print(f"Erro no treino: {str(e)}")
|
|
with _status_lock:
|
|
_treinamento_status = {
|
|
"status": "erro",
|
|
"ambiente": ambiente,
|
|
"versao": None,
|
|
"detalhe": str(e)[:500],
|
|
}
|
|
|
|
|
|
@app.post("/detectar")
|
|
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:
|
|
Image.MAX_IMAGE_PIXELS = None
|
|
modelo = carregar_modelo_do_s3(ambiente)
|
|
conteudo = await file.read()
|
|
imagem = Image.open(io.BytesIO(conteudo))
|
|
results = modelo(imagem, conf=0.25)
|
|
deteccoes = [
|
|
{"box": [round(x, 2) for x in b.xyxy[0].tolist()], "conf": round(float(b.conf), 2), "class": int(b.cls)}
|
|
for r in results for b in r.boxes
|
|
]
|
|
return {"status": "sucesso", "deteccoes": deteccoes}
|
|
except Exception as e:
|
|
log_print(f"Erro deteccao: {str(e)}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.post("/treinar")
|
|
async def treinar(dados: dict):
|
|
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:
|
|
if _treinamento_status["status"] == "running":
|
|
raise HTTPException(status_code=409, detail="Treinamento já em andamento")
|
|
pular_triagem = dados.get("pular_triagem", False)
|
|
forcar_do_zero = dados.get("forcar_do_zero", False)
|
|
_treinamento_status = {"status": "running", "ambiente": ambiente, "versao": None, "detalhe": None}
|
|
|
|
threading.Thread(target=_treinar_bg, args=(ambiente, pular_triagem, forcar_do_zero), daemon=True).start()
|
|
return {"status": "iniciado", "ambiente": ambiente}
|
|
|
|
|
|
@app.get("/treinar/status")
|
|
async def status_treino():
|
|
with _status_lock:
|
|
return {**dict(_treinamento_status), "logs": list(_training_logs[-30:])}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|