Update app.py

app.py

@@ -1,37 +1,34 @@
 # app.py
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
-import numpy as np
-import cv2
-import base64
-import io
 import os
 import uuid
 import boto3
-from supabase import create_client, Client
-
-# Variaveis de ambiente
+import supabase
+import numpy as np
+import cv2
+import base64
+import mimetypes
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from botocore.exceptions import NoCredentialsError
 
-# --- AWS S3 ---
-AWS_ACCESS_KEY_ID = os.environ.get("AWS_ACCESS_KEY_ID")
-AWS_SECRET_ACCESS_KEY = os.environ.get("AWS_SECRET_ACCESS_KEY")
-AWS_S3_BUCKET_NAME = os.environ.get("AWS_S3_BUCKET_NAME")
-AWS_S3_REGION = os.environ.get("AWS_S3_REGION")
+# ==============================================================================
+# 1. CONFIGURAÇÃO E CLIENTES (CARREGADOS NA INICIALIZAÇÃO)
+# ==============================================================================
 
-# --- Supabase ---
-SUPABASE_URL = os.environ.get("SUPABASE_URL")
-SUPABASE_KEY = os.environ.get("SUPABASE_KEY")
+# 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")
+AWS_S3_BUCKET_NAME = os.getenv("AWS_S3_BUCKET_NAME")
+AWS_S3_REGION = os.getenv("AWS_S3_REGION")
+SUPABASE_URL = os.getenv("SUPABASE_URL")
+SUPABASE_KEY = os.getenv("SUPABASE_KEY")
 
-# Supabase
-try:
-    supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)
-    print("Cliente Supabase inicializado com sucesso.")
-except Exception as e:
-    print(f"ERRO ao inicializar Supabase: {e}")
-    raise
+# Validação para garantir que todos os secrets foram configurados
+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.")
 
-# S3 (boto3)
+# Inicializa os clientes globais
 try:
     s3_client = boto3.client(
         's3',
@@ -39,14 +36,15 @@ try:
         aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
         region_name=AWS_S3_REGION
     )
-    print("Cliente S3 (boto3) inicializado com sucesso.")
+    supabase_client = supabase.create_client(SUPABASE_URL, SUPABASE_KEY)
+    print("Clientes S3 e Supabase inicializados com sucesso.")
 except Exception as e:
-    print(f"ERRO ao inicializar S3: {e}")
-    raise
+    raise RuntimeError(f"Erro ao inicializar os clientes: {e}")
 
-app = FastAPI(title="Detector de Corrosão Branca")
+# Inicializa a aplicação FastAPI
+app = FastAPI(title="Detector de Corrosão Branca com Salvamento S3/Supabase")
 
-# PARA PROTOTIPO: permitir todas origens. Em produção restrinja ao domínio do frontend.
+# Em produção, restrinja ao domínio do seu frontend
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -55,155 +53,179 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# ==============================================================================
+# 2. FUNÇÕES HELPER (LÓGICA REUTILIZÁVEL)
+# ==============================================================================
+
 def process_image_bytes(img_bytes: bytes):
-    # lê bytes em numpy + OpenCV
+    """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.")
 
-    # 1) Converter para HSV
     hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
-
-    # 2) máscara do fundo preto (V baixo)
     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)
-
-    # 3) objeto = invertendo máscara do fundo
     mask_obj = cv2.bitwise_not(mask_bg)
-
-    # 4) limpar máscara (morfologia)
     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)
-
-    # 5) maior contorno (supõe um parafuso)
     contours, _ = cv2.findContours(mask_obj, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if not contours:
-        return {"error": "Nenhum objeto detectado"}
+        raise ValueError("Nenhum objeto principal detectado na imagem.")
 
     largest = max(contours, key=cv2.contourArea)
     mask_clean = np.zeros_like(mask_obj)
     cv2.drawContours(mask_clean, [largest], -1, 255, cv2.FILLED)
-
-    # 6) isolar objeto
     isolated = cv2.bitwise_and(img, img, mask=mask_clean)
-
-    # 7) detectar corrosão BRANCA (S baixa, V alta)
     hsv_iso = cv2.cvtColor(isolated, cv2.COLOR_BGR2HSV)
     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)
     mask_white = cv2.bitwise_and(mask_white, mask_white, mask=mask_clean)
 
-    # 8) métricas
     total_pixels = int(np.count_nonzero(mask_clean))
     corrosion_pixels = int(np.count_nonzero(mask_white))
     percent = (corrosion_pixels / max(1, total_pixels)) * 100.0
 
-    # 9) preparar imagens para frontend (PNG base64)
-    # isolado em RGB para visualização
     isolated_rgb = cv2.cvtColor(isolated, cv2.COLOR_BGR2RGB)
-    corrosion_vis = cv2.bitwise_and(isolated_rgb, isolated_rgb, mask=mask_white)
+    corrosion_vis_rgb = cv2.bitwise_and(isolated_rgb, isolated_rgb, mask=mask_white)
 
-    def to_data_uri(img_arr):
-        # img_arr: RGB uint8
-        bgr = cv2.cvtColor(img_arr, cv2.COLOR_RGB2BGR)
+    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}"
 
-    isolated_b64 = to_data_uri(isolated_rgb)
-    corrosion_b64 = to_data_uri(corrosion_vis)
-
-    return {
+    analysis_results = {
         "percent": round(percent, 4),
         "total_pixels": total_pixels,
         "corrosion_pixels": corrosion_pixels,
-        "isolated_image": isolated_b64,
-        "corrosion_image": corrosion_b64,
+        "isolated_image": to_data_uri(isolated_rgb),
+        "corrosion_image": to_data_uri(corrosion_vis_rgb),
     }
+    
+    return analysis_results, corrosion_vis_rgb
+
+# ==============================================================================
+# 3. ENDPOINTS DA API
+# ==============================================================================
+
+@app.get("/")
+def read_root():
+    return {"status": "ok", "message": "API de Análise de Corrosão no ar!"}
 
-def image_link_gen(amostra_id: int) -> dict | None:
+@app.post("/analyze")
+async def analyze(file: UploadFile = File(...)):
     """
-    Busca uma amostra pelo ID no Supabase e gera URLs pré-assinadas para suas imagens.
-    Retorna um dicionário com os links ou None se a amostra não for encontrada.
+    Endpoint principal: Analisa a imagem, salva original e resultado no S3,
+    e persiste os dados no Supabase.
     """
-    print(f"\nBuscando amostra com ID: {amostra_id} no Supabase...")
+    content = await file.read()
+    s3_key_original = None
+    s3_key_resultado = None
+
+    try:
+        # 1. Processa a imagem para obter métricas e a imagem de resultado
+        analysis_results, corrosion_image_np = process_image_bytes(content)
+
+        # 2. Faz o upload da IMAGEM ORIGINAL para o S3
+        content_type = file.content_type or 'application/octet-stream'
+        extensao = mimetypes.guess_extension(content_type) or '.jpg'
+        s3_key_original = f"imagens_originais/{uuid.uuid4()}{extensao}"
         
-    # 1. Busca os dados da amostra no banco
-    response = supabase.from_("amostra").select("imagem_original, imagem_resultado").eq("id", amostra_id).single().execute()
+        print(f"Fazendo upload da imagem original para: {s3_key_original}")
+        s3_client.put_object(
+            Bucket=AWS_S3_BUCKET_NAME,
+            Key=s3_key_original,
+            Body=content,
+            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"
         
-    amostra = response.data
-    if not amostra:
-        print(f"Amostra com ID {amostra_id} não encontrada.")
-        return None
+        print(f"Fazendo upload da imagem de resultado para: {s3_key_resultado}")
+        bgr_result = cv2.cvtColor(corrosion_image_np, cv2.COLOR_RGB2BGR)
+        ok, buffer = cv2.imencode('.png', bgr_result)
+        if not ok:
+            raise ValueError("Falha ao codificar a imagem de resultado para PNG.")
         
-    print("Amostra encontrada. Gerando links temporários...")
+        s3_client.put_object(
+            Bucket=AWS_S3_BUCKET_NAME,
+            Key=s3_key_resultado,
+            Body=buffer.tobytes(),
+            ContentType='image/png'
+        )
+
+        # 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
+        }
+        response = supabase_client.from_("amostras").insert(dados_para_inserir).execute()
         
-    # Pega as chaves S3 do resultado
-    s3_key_original = amostra.get("imagem_original")
-    s3_key_resultado = amostra.get("imagem_resultado")
+        # Adiciona o ID do novo registro à resposta
+        new_record_id = response.data[0]['id']
+        analysis_results['database_id'] = new_record_id
         
-    links = {}
+        return JSONResponse(content=analysis_results)
+
+    except Exception as e:
+        # Lógica de Rollback: Se algo falhar, tenta limpar os arquivos do S3
+        print(f"ERRO no processo de análise: {e}")
+        if s3_key_original:
+            print(f"Removendo objeto órfão do S3: {s3_key_original}")
+            s3_client.delete_object(Bucket=AWS_S3_BUCKET_NAME, Key=s3_key_original)
+        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)
         
-    # 2. Gera URL pré-assinada para a imagem original
-    if s3_key_original:
-        url_original = s3_client.generate_presigned_url(
-            'get_object',
-            Params={'Bucket': AWS_S3_BUCKET_NAME, 'Key': s3_key_original},
-            ExpiresIn=3600  # Link válido por 1 hora (3600 segundos)
-        )
-        links['original'] = url_original
-            
-    # 3. Gera URL pré-assinada para a imagem de resultado (se existir)
-    if s3_key_resultado:
-        url_resultado = s3_client.generate_presigned_url(
-            'get_object',
-            Params={'Bucket': AWS_S3_BUCKET_NAME, 'Key': s3_key_resultado},
-            ExpiresIn=3600
-        )
-        links['resultado'] = url_resultado
-            
-    return links
+        raise HTTPException(status_code=500, detail=f"Ocorreu um erro interno: {e}")
 
-@app.post("/analyze")
-async def analyze(file: UploadFile = File(...)):
-    content = await file.read()
-    try:
-        result = process_image_bytes(content)
-    except ValueError as e:
-        raise HTTPException(status_code=400, detail=str(e))
-    return JSONResponse(result)
 
-@app.get("/links/{amostra_id}")
-async def get_link(amostra_id: int):
+@app.get("/samples/{sample_id}")
+async def get_sample_images(sample_id: int):
+    """
+    Endpoint de recuperação: Busca uma amostra pelo ID e retorna
+    links temporários (pré-assinados) para suas imagens no S3.
+    """
     try:
+        print(f"Buscando amostra com ID: {sample_id}...")
+        response = supabase_client.from_("amostras").select("imagem_original, imagem_resultado").eq("id", sample_id).single().execute()
         
-        links = image_link_gen(amostra_id)
+        if not response.data:
+            raise HTTPException(status_code=404, detail=f"Amostra com ID {sample_id} não encontrada.")
+            
+        amostra = response.data
+        s3_key_original = amostra.get("imagem_original")
+        s3_key_resultado = amostra.get("imagem_resultado")
+        
+        links = {}
         
-        if links is None:
-            raise HTTPException(
-                status_code=404, 
-                detail=f"Amostra com ID {amostra_id} não encontrada."
+        if s3_key_original:
+            links['url_original'] = s3_client.generate_presigned_url(
+                'get_object',
+                Params={'Bucket': AWS_S3_BUCKET_NAME, 'Key': s3_key_original},
+                ExpiresIn=3600 # 1 hora
             )
         
-        return links
+        if s3_key_resultado:
+            links['url_resultado'] = s3_client.generate_presigned_url(
+                'get_object',
+                Params={'Bucket': AWS_S3_BUCKET_NAME, 'Key': s3_key_resultado},
+                ExpiresIn=3600 # 1 hora
+            )
+            
+        return JSONResponse(content=links)
 
     except Exception as e:
-
-        print(f"Erro inesperado ao processar ID {amostra_id}: {e}")
-        
-        raise HTTPException(
-            status_code=500, 
-            detail=f"Erro interno no servidor ao tentar buscar os links. {e}"
-        )
-    
-
-@app.get("/")
-def read_root():
-    return {"status": "ok"}
-
+        raise HTTPException(status_code=500, detail=f"Ocorreu um erro interno ao buscar os links: {e}")