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') 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: """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 cerebro.pt atual só para efeito de treino — o arquivo em # modelos/{ambiente}/atual/cerebro.pt não é tocado, então o /detectar # em produção continua servindo normalmente enquanto este treino roda. 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. 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) else: modelo_base.train(data=yaml_path, epochs=30, imgsz=640, batch=8, device='cpu', plots=True, lr0=0.01, lrf=0.01) modelo_base.reset_callbacks() 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 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(...)): 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 with _status_lock: if _treinamento_status["status"] == "running": raise HTTPException(status_code=409, detail="Treinamento já em andamento") ambiente = dados.get("ambiente", "gondola") 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)