Update app.py

app.py

@@ -7,6 +7,12 @@ import numpy as np
 import cv2
 import base64
 import mimetypes
+# --- NOVAS importações do 'teste yolo.txt' ---
+import io
+from pathlib import Path
+from PIL import Image
+from ultralytics import YOLO
+# --- Fim das novas importações ---
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
@@ -16,6 +22,11 @@ from botocore.exceptions import NoCredentialsError
 # 1. CONFIGURAÇÃO E CLIENTES (CARREGADOS NA INICIALIZAÇÃO)
 # ==============================================================================
 
+# --- NOVAS constantes do 'teste yolo.txt' ---
+LOCAL_RESULT = 'back-ciser/result'
+YOLO_MODEL_PATH = 'back-ciser/best.pt' # Caminho para o seu modelo treinado
+# --- Fim das novas constantes ---
+
 # Carrega as configurações das variáveis de ambiente (Secrets do Hugging Face)
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -24,7 +35,7 @@ AWS_S3_REGION = os.getenv("AWS_S3_REGION")
 SUPABASE_URL = os.getenv("SUPABASE_URL")
 SUPABASE_KEY = os.getenv("SUPABASE_KEY")
 
-# Validação para garantir que todos os secrets foram configurados
+# Validação (do back atual.txt)
 if not all([AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_S3_BUCKET_NAME, AWS_S3_REGION, SUPABASE_URL, SUPABASE_KEY]):
     raise RuntimeError("Erro: Nem todos os secrets necessários foram configurados no Hugging Face Space.")
 
@@ -44,7 +55,7 @@ except Exception as e:
 # Inicializa a aplicação FastAPI
 app = FastAPI(title="Detector de Corrosão Branca com Salvamento S3/Supabase")
 
-# Em produção, restrinja ao domínio do seu frontend
+# Middleware (do back atual.txt)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -57,25 +68,23 @@ app.add_middleware(
 # 2. FUNÇÕES HELPER (LÓGICA REUTILIZÁVEL)
 # ==============================================================================
 
+# --- Função ORIGINAL do 'back atual.txt' (usada por ambas as lógicas) ---
 def process_image_bytes(img_bytes: bytes):
-    """Lógica de análise de imagem com OpenCV. Retorna um dict com métricas e a imagem de corrosão em numpy (RGB)."""
+    """Lógica de análise de imagem com OpenCV.
+    Retorna um dict com métricas e a imagem de corrosão em numpy."""
     nparr = np.frombuffer(img_bytes, np.uint8)
     img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
     if img is None:
         raise ValueError("Não foi possível decodificar a imagem.")
 
     hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
-
-    # Segmentação do objeto principal (fundo escuro)
     lower_bg = np.array([0, 0, 0], dtype=np.uint8)
     upper_bg = np.array([180, 255, 50], dtype=np.uint8)
     mask_bg = cv2.inRange(hsv, lower_bg, upper_bg)
     mask_obj = cv2.bitwise_not(mask_bg)
-
     kernel = np.ones((5, 5), np.uint8)
     mask_obj = cv2.morphologyEx(mask_obj, cv2.MORPH_OPEN, kernel)
     mask_obj = cv2.morphologyEx(mask_obj, cv2.MORPH_CLOSE, kernel)
-
     contours, _ = cv2.findContours(mask_obj, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if not contours:
         raise ValueError("Nenhum objeto principal detectado na imagem.")
@@ -83,11 +92,8 @@ def process_image_bytes(img_bytes: bytes):
     largest = max(contours, key=cv2.contourArea)
     mask_clean = np.zeros_like(mask_obj)
     cv2.drawContours(mask_clean, [largest], -1, 255, cv2.FILLED)
-
     isolated = cv2.bitwise_and(img, img, mask=mask_clean)
     hsv_iso = cv2.cvtColor(isolated, cv2.COLOR_BGR2HSV)
-
-    # Corrosão branca (faixa de brancos)
     lower_white = np.array([0, 0, 180], dtype=np.uint8)
     upper_white = np.array([180, 60, 255], dtype=np.uint8)
     mask_white = cv2.inRange(hsv_iso, lower_white, upper_white)
@@ -103,8 +109,7 @@ def process_image_bytes(img_bytes: bytes):
     def to_data_uri(img_arr_rgb):
         bgr = cv2.cvtColor(img_arr_rgb, cv2.COLOR_RGB2BGR)
         ok, buf = cv2.imencode(".png", bgr)
-        if not ok:
-            return None
+        if not ok: return None
         b64 = base64.b64encode(buf.tobytes()).decode("ascii")
         return f"data:image/png;base64,{b64}"
 
@@ -115,14 +120,17 @@ def process_image_bytes(img_bytes: bytes):
         "isolated_image": to_data_uri(isolated_rgb),
         "corrosion_image": to_data_uri(corrosion_vis_rgb),
     }
-
+    
     return analysis_results, corrosion_vis_rgb
 
-    
+
+# --- NOVAS funções helper do 'teste yolo.txt' (para o fluxo de segmentação) ---
+
 def analyze_and_persist(img_bytes: bytes, logical_name: str):
     """
     Processa a imagem em memória (process_image_bytes), faz upload no S3 (original & resultado)
-    e insere na tabela 'amostras'. Retorna o dict da análise + chaves S3 + database_id.
+    e insere na tabela 'amostras'.
+    Retorna o dict da análise + chaves S3 + database_id.
     """
     # 1) Processa
     analysis_results, corrosion_image_np = process_image_bytes(img_bytes)
@@ -152,7 +160,7 @@ def analyze_and_persist(img_bytes: bytes, logical_name: str):
         ContentType='image/png'
     )
 
-    # 4) Insert Supabase (na mesma tabela 'amostras')
+    # 4) Insert Supabase (na mesma tabela 'amostra')
     dados_para_inserir = {
         "imagem_original": s3_key_original,
         "imagem_resultado": s3_key_resultado,
@@ -170,15 +178,16 @@ def analyze_and_persist(img_bytes: bytes, logical_name: str):
     }
     return out
 
-    def segmentate(file_bytes: bytes):
+
+def segmentate(file_bytes: bytes):
     """
     Roda o YOLO, salva predição e crops, e retorna:
     - save_dir (Path): diretório raiz onde o Ultralytics salvou esta execução
     - crop_paths (list[Path]): lista com todos os arquivos de imagem dentro de /crops
     """
-    modelo = YOLO('back-ciser/best.pt')
+    modelo = YOLO(YOLO_MODEL_PATH) # Usando a constante definida no início
 
-    # Passa a imagem como array (sem salvar a original só para inferir)
+    # Passa a imagem como array
     img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
     np_img = np.array(img)
     np_img_bgr = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)
@@ -191,7 +200,7 @@ def analyze_and_persist(img_bytes: bytes, logical_name: str):
         save=True,
         save_crop=True,
         project=LOCAL_RESULT,
-        name="predict",      # deixa previsível
+        name="predict",
         exist_ok=True,
         verbose=False
     )
@@ -203,7 +212,7 @@ def analyze_and_persist(img_bytes: bytes, logical_name: str):
     if not crops_dir.exists():
         return save_dir, []
 
-    # Varre recursivamente (crops/<classe>/*.jpg|png|...)
+    # Varre recursivamente
     crop_paths = []
     for ext in ("*.jpg", "*.png", "*.jpeg", "*.bmp", "*.tif", "*.tiff"):
         crop_paths.extend(crops_dir.rglob(ext))
@@ -211,6 +220,9 @@ def analyze_and_persist(img_bytes: bytes, logical_name: str):
     crop_paths = sorted(crop_paths, key=lambda p: str(p).lower())
     return save_dir, crop_paths
 
+# --- Fim das novas funções helper ---
+
+
 # ==============================================================================
 # 3. ENDPOINTS DA API
 # ==============================================================================
@@ -219,11 +231,13 @@ def analyze_and_persist(img_bytes: bytes, logical_name: str):
 def read_root():
     return {"status": "ok", "message": "API de Análise de Corrosão no ar!"}
 
+
+# --- Endpoint ORIGINAL do 'back atual.txt' (Sem alterações) ---
 @app.post("/analyze")
 async def analyze(file: UploadFile = File(...)):
     """
     Endpoint principal: Analisa a imagem, salva original e resultado no S3,
-    e persiste os dados no Supabase.
+    e persiste os dados no Supabase. (LÓGICA ANTIGA)
     """
     content = await file.read()
     s3_key_original = None
@@ -246,6 +260,7 @@ async def analyze(file: UploadFile = File(...)):
             ContentType=content_type
         )
 
+        
         # 3. Faz o upload da IMAGEM DE RESULTADO (em memória) para o S3
         s3_key_resultado = f"imagens_resultados/{uuid.uuid4()}.png"
         
@@ -265,10 +280,6 @@ async def analyze(file: UploadFile = File(...)):
         # 4. Salva os metadados e as chaves S3 no Supabase
         print("Inserindo registro no Supabase...")
         dados_para_inserir = {
-            #"nome_amostra": file.filename,
-            #"percentual_corrosao": analysis_results["percent"],
-            #"pixels_totais_obj": analysis_results["total_pixels"],
-            #"pixels_corrosao": analysis_results["corrosion_pixels"],
             "imagem_original": s3_key_original,
             "imagem_resultado": s3_key_resultado,
             "resultado": analysis_results["percent"]
@@ -282,7 +293,7 @@ async def analyze(file: UploadFile = File(...)):
         return JSONResponse(content=analysis_results)
 
     except Exception as e:
-        # Lógica de Rollback: Se algo falhar, tenta limpar os arquivos do S3
+        # Lógica de Rollback:
         print(f"ERRO no processo de análise: {e}")
         if s3_key_original:
             print(f"Removendo objeto órfão do S3: {s3_key_original}")
@@ -290,10 +301,59 @@ async def analyze(file: UploadFile = File(...)):
         if s3_key_resultado:
             print(f"Removendo objeto órfão do S3: {s3_key_resultado}")
             s3_client.delete_object(Bucket=AWS_S3_BUCKET_NAME, Key=s3_key_resultado)
-        
+ 
+        raise HTTPException(status_code=500, detail=f"Ocorreu um erro interno: {e}")
+
+
+# --- NOVO Endpoint (lógica do 'teste yolo.txt') ---
+@app.post("/segment_and_analyze")
+async def segment_and_analyze(file: UploadFile = File(...)):
+    """
+    Endpoint com YOLO: Segmenta a imagem, analisa a original E cada crop,
+    e salva tudo no S3/Supabase.
+    (Lógica dos sources 36-39)
+    """
+    content = await file.read()
+
+    # Para rollback em caso de erro
+    created_s3_keys = []
+
+    try:
+        # 0) Segmenta e obtém caminhos dos crops
+        save_dir, crop_paths = segmentate(content)
+
+        # 1) Analisa a imagem ORIGINAL
+        original_res = analyze_and_persist(content, logical_name=file.filename or "upload_sem_nome")
+        created_s3_keys.extend([original_res["s3_original"], original_res["s3_resultado"]])
+
+        # 2) Analisa cada CROP
+        crop_analyses = []
+        for crop_path in crop_paths:
+            crop_bytes = crop_path.read_bytes()
+            logical_name = str(crop_path.relative_to(save_dir)).replace("\\", "/") # nome amigável
+            crop_res = analyze_and_persist(crop_bytes, logical_name=logical_name)
+            created_s3_keys.extend([crop_res["s3_original"], crop_res["s3_resultado"]])
+            crop_res["crop_path"] = logical_name
+            crop_analyses.append(crop_res)
+
+        # 3) Resposta consolidada
+        return JSONResponse(content={
+            "original": original_res,
+            "crops_count": len(crop_analyses),
+            "crops": crop_analyses
+        })
+
+    except Exception as e:
+        # Rollback simples: tenta apagar do S3 o que já foi criado
+        for key in created_s3_keys:
+            try:
+                s3_client.delete_object(Bucket=AWS_S3_BUCKET_NAME, Key=key)
+            except Exception:
+                pass
         raise HTTPException(status_code=500, detail=f"Ocorreu um erro interno: {e}")
 
 
+# --- Endpoint ORIGINAL do 'back atual.txt' (Sem alterações) ---
 @app.get("/samples/{sample_id}")
 async def get_sample_images(sample_id: int):
     """