app2.py

# app2.py
import os
import uuid
import boto3
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
from ultralytics import YOLO

# ==============================================================================
# 1. CONFIGURAÇÃO E CLIENTES (CARREGADOS NA INICIALIZAÇÃO)
# ==============================================================================

# 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")

# 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.")

# Inicializa os clientes globais
try:
    s3_client = boto3.client(
        's3',
        aws_access_key_id=AWS_ACCESS_KEY_ID,
        aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
        region_name=AWS_S3_REGION
    )
    supabase_client = supabase.create_client(SUPABASE_URL, SUPABASE_KEY)
    print("Clientes S3 e Supabase inicializados com sucesso.")
except Exception as e:
    raise RuntimeError(f"Erro ao inicializar os clientes: {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
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ==============================================================================
# 2. FUNÇÕES HELPER (LÓGICA REUTILIZÁVEL)
# ==============================================================================

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."""
    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)
    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.")

    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)
    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)

    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

    isolated_rgb = cv2.cvtColor(isolated, cv2.COLOR_BGR2RGB)
    corrosion_vis_rgb = cv2.bitwise_and(isolated_rgb, isolated_rgb, mask=mask_white)

    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
        b64 = base64.b64encode(buf.tobytes()).decode("ascii")
        return f"data:image/png;base64,{b64}"

    analysis_results = {
        "percent": round(percent, 4),
        "total_pixels": total_pixels,
        "corrosion_pixels": corrosion_pixels,
        "isolated_image": to_data_uri(isolated_rgb),
        "corrosion_image": to_data_uri(corrosion_vis_rgb),
    }
    
    return analysis_results, corrosion_vis_rgb

def segmentate():

    # Predição basead no melhor modelo gerado (best.pt)
    modelo = YOLO('/content/resultado/Parafuso3/weights/best.pt')
    resultados = modelo.predict(
    source='/content/dados/images/train',
    imgsz=1280,
    conf=0.7,
    iou=0.3,
    save=True,
    save_crop=True,
    project='/content/resultado'
)

    
    

# ==============================================================================
# 3. ENDPOINTS DA API
# ==============================================================================

@app.get("/")
def read_root():
    return {"status": "ok", "message": "API de Análise de Corrosão no ar!"}


@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.
    """
    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}"
        
        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"
        
        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.")
        
        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()
        
        # Adiciona o ID do novo registro à resposta
        new_record_id = response.data[0]['id']
        analysis_results['database_id'] = new_record_id
        
        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)
        
        raise HTTPException(status_code=500, detail=f"Ocorreu um erro interno: {e}")


@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()
        
        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 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
            )
        
        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:
        raise HTTPException(status_code=500, detail=f"Ocorreu um erro interno ao buscar os links: {e}")