from fastapi import FastAPI, File, UploadFile, Form, HTTPException from ultralytics import YOLO import boto3 import os import io import shutil import yaml 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') MODELOS_CARREGADOS = {} 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") print(f"[{carimbo}] {msg}", flush=True) 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 cérebro '{ambiente}'...") s3.download_file(BUCKET, s3_key, local_path) MODELOS_CARREGADOS[ambiente] = YOLO(local_path) return MODELOS_CARREGADOS[ambiente] @app.post("/detectar") async def detectar(ambiente: str = Form(...), file: UploadFile = File(...)): 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 detecção: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) @app.post("/treinar") async def treinar(dados: dict): ambiente = dados.get("ambiente", "gondola") pular_triagem = dados.get("pular_triagem", False) try: prefix_novos = f"treinamento/{ambiente}/novos-treinamentos/" objs = s3.list_objects_v2(Bucket=BUCKET, Prefix=prefix_novos) modelo_base = carregar_modelo_do_s3(ambiente) 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 # RÉGUA AJUSTADA PARA 0.30 PARA ACEITAR MAIS IMAGENS decisao = "base-oficial" if media >= 0.30 or media == 0 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: continue dataset_local = f"/tmp/dataset_{ambiente}" img_dir, lbl_dir = f"{dataset_local}/train/images", 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) if len(os.listdir(img_dir)) > 0: yaml_path = f"{dataset_local}/data.yaml" 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. Foco no mAP50...") # O YOLO imprimirá o mAP automaticamente no log do Docker aqui: modelo_base.train(data=yaml_path, epochs=30, imgsz=640, batch=16, device='cpu', plots=True) best = "runs/detect/train/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 CONCLUÍDO!") return {"status": "sucesso", "versao": carimbo} else: log_print("⚠️ Nenhuma imagem passou na triagem para a Base Oficial.") return {"status": "vazio"} except Exception as e: log_print(f"💥 Erro: {str(e)}") return {"status": "erro", "detalhe": str(e)} if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)