app3.py

import os
import uuid
import json
import mimetypes
import base64
from typing import List, Tuple, Dict, Any

import boto3
import supabase
import numpy as np
import cv2

from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from botocore.exceptions import NoCredentialsError
from ultralytics import YOLO

# ==============================================================================
# 1. CONFIGURAÇÃO
# ==============================================================================

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")
YOLO_MODEL_PATH = os.getenv("YOLO_MODEL_PATH", "best.pt")

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: faltam secrets da AWS ou Supabase.")

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.")
except Exception as e:
    raise RuntimeError(f"Erro ao inicializar clientes: {e}")

try:
    yolo_model = YOLO(YOLO_MODEL_PATH)
    print(f"YOLO carregado: {YOLO_MODEL_PATH}")
except Exception as e:
    raise RuntimeError(f"Falha ao carregar YOLO: {e}")

app = FastAPI(title="CorroScan API — YOLO + OpenCV")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # restrinja em produção
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ==============================================================================
# 2. HELPERS
# ==============================================================================



def _to_data_uri_from_rgb(img_rgb: np.ndarray) -> str:
    if img_rgb is None or img_rgb.size == 0:
        return None
    bgr = cv2.cvtColor(img_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}"

def draw_boxes_on_bgr(img_bgr: np.ndarray, boxes_xyxy: np.ndarray, labels: List[str]) -> np.ndarray:
    out = img_bgr.copy()
    for (x1, y1, x2, y2), label in zip(boxes_xyxy, labels):
        x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
        cv2.rectangle(out, (x1, y1), (x2, y2), (0, 200, 0), 2)
        cv2.putText(out, label, (x1, max(y1 - 5, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 220, 0), 2, cv2.LINE_AA)
    return cv2.cvtColor(out, cv2.COLOR_BGR2RGB)

def circular_roi_from_mask(mask_clean: np.ndarray, shrink: float = 0.9) -> np.ndarray:
    cnts, _ = cv2.findContours(mask_clean, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not cnts:
        return mask_clean.copy()
    c = max(cnts, key=cv2.contourArea)
    (x, y), r = cv2.minEnclosingCircle(c)
    r = max(1, int(r * shrink))
    cx, cy = int(x), int(y)
    roi = np.zeros_like(mask_clean)
    cv2.circle(roi, (cx, cy), r, 255, -1)
    return roi

def remove_specular_highlights(hsv_iso: np.ndarray, mask: np.ndarray, v_spec: int = 230) -> np.ndarray:
    v = hsv_iso[..., 2]
    spec = cv2.inRange(v, v_spec, 255)
    return cv2.bitwise_and(mask, cv2.bitwise_not(spec))

def remove_small_components(mask: np.ndarray, min_area_ratio: float, also_border: bool = True) -> np.ndarray:
    if mask.max() == 0:
        return mask
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
    total = int(np.count_nonzero(mask))
    min_area = max(1, int(total * min_area_ratio))
    out = np.zeros_like(mask)
    H, W = mask.shape[:2]
    for i in range(1, num_labels):
        x, y, w, h, area = stats[i]
        if area < min_area:
            continue
        if also_border and (x == 0 or y == 0 or x + w == W or y + h == H):
            continue
        out[labels == i] = 255
    return out

def dilate_around(mask: np.ndarray, it: int = 1) -> np.ndarray:
    k = np.ones((3, 3), np.uint8)
    return cv2.dilate(mask, k, iterations=it)

def illum_normalize_L(img_bgr: np.ndarray, sigma: float = 21) -> np.ndarray:
    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
    L = lab[..., 0].astype(np.float32)
    base = cv2.GaussianBlur(L, (0, 0), sigma)
    base = np.maximum(base, 1.0)
    Ln = (L / base) * 128.0
    Ln = np.clip(Ln, 0, 255).astype(np.uint8)
    return Ln

def texture_map(Ln: np.ndarray, ksize: int = 3) -> np.ndarray:
    lap = cv2.Laplacian(Ln, cv2.CV_16S, ksize=ksize)
    t = np.abs(lap).astype(np.uint16)
    t = np.clip(t, 0, 255).astype(np.uint8)
    return t

def get_corrosion_mask(hsv_iso: np.ndarray, mask_clean: np.ndarray, mode: str, src_bgr_iso: np.ndarray = None) -> np.ndarray:
    mode = (mode or "white").lower()
    kernel = np.ones((3, 3), np.uint8)

    if mode == "white":
        # cor conservadora
        lower = np.array([0, 0, 115], dtype=np.uint8)
        upper = np.array([180, 45, 215], dtype=np.uint8)
        m_color = cv2.inRange(hsv_iso, lower, upper)

        # especular duro e vizinhança
        v = hsv_iso[..., 2]
        s = hsv_iso[..., 1]
        spec_core = cv2.inRange(v, 220, 255) & cv2.inRange(s, 0, 40)
        spec = dilate_around(spec_core, it=2)

        if src_bgr_iso is None:
            raise ValueError("src_bgr_iso é necessário para textura no modo white.")

        Ln = illum_normalize_L(src_bgr_iso)
        tmap = texture_map(Ln, ksize=3)
        tmask = cv2.inRange(tmap, 8, 255)

        m = m_color
        m = cv2.bitwise_and(m, cv2.bitwise_not(spec))
        m = cv2.bitwise_and(m, tmask)

    elif mode == "black":
        lower = np.array([0, 0, 0], dtype=np.uint8)
        upper = np.array([180, 255, 60], dtype=np.uint8)
        m = cv2.inRange(hsv_iso, lower, upper)

    elif mode == "red":
        lower1 = np.array([0, 80, 60], dtype=np.uint8)
        upper1 = np.array([10, 255, 255], dtype=np.uint8)
        lower2 = np.array([170, 80, 60], dtype=np.uint8)
        upper2 = np.array([180, 255, 255], dtype=np.uint8)
        m = cv2.bitwise_or(cv2.inRange(hsv_iso, lower1, upper1),
                           cv2.inRange(hsv_iso, lower2, upper2))
    else:
        return get_corrosion_mask(hsv_iso, mask_clean, "white", src_bgr_iso)

    m = cv2.bitwise_and(m, m, mask=mask_clean)
    m = cv2.morphologyEx(m, cv2.MORPH_OPEN, kernel, iterations=1)
    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, kernel, iterations=1)
    m = remove_small_components(m, min_area_ratio=0.005, also_border=True)
    return m

def process_image_bgr(img_bgr: np.ndarray, corrosion_type: str = "white") -> Tuple[Dict[str, Any], np.ndarray]:
    if img_bgr is None or img_bgr.size == 0:
        raise ValueError("Imagem vazia.")

    # suaviza reflexos preservando bordas
    img_bgr = cv2.bilateralFilter(img_bgr, d=7, sigmaColor=60, sigmaSpace=60)
    hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)

    # objeto principal
    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 detectado no crop.")

    largest = max(contours, key=cv2.contourArea)
    mask_clean = np.zeros_like(mask_obj)
    cv2.drawContours(mask_clean, [largest], -1, 255, cv2.FILLED)

    # foca no tampo
    roi_circle = circular_roi_from_mask(mask_clean, shrink=0.9)
    mask_clean = cv2.bitwise_and(mask_clean, roi_circle)

    isolated = cv2.bitwise_and(img_bgr, img_bgr, mask=mask_clean)
    hsv_iso = cv2.cvtColor(isolated, cv2.COLOR_BGR2HSV)

    mask_corrosion = get_corrosion_mask(
        hsv_iso=hsv_iso,
        mask_clean=mask_clean,
        mode=corrosion_type,
        src_bgr_iso=isolated
    )

    total_pixels = int(np.count_nonzero(mask_clean))
    corrosion_pixels = int(np.count_nonzero(mask_corrosion))
    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_corrosion)

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

# ==============================================================================
# 3. ENDPOINTS
# ==============================================================================

@app.get("/")
def read_root():
    return {"status": "ok", "message": "API YOLO + OpenCV pronta."}

@app.post("/analyze")
async def analyze(
    file: UploadFile = File(...),
    corrosion_type: str = Form("white")  # "white" | "black" | "red"
):
    """
    1) YOLO detecta parafusos e gera boxes
    2) Para cada box, recorta e roda OpenCV conforme corrosion_type
    3) Sobe original e crops no S3
    4) Insere metadados básicos no Supabase
    """
    content = await file.read()
    s3_key_original = None
    s3_keys_crops: List[str] = []

    try:
        # decodifica
        nparr = np.frombuffer(content, np.uint8)
        img_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        if img_bgr is None:
            raise ValueError("Imagem inválida.")

        # YOLO
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        yolo_results = yolo_model.predict(source=img_rgb, imgsz=1280, conf=0.5, iou=0.4, verbose=False)
        if not yolo_results:
            raise ValueError("YOLO não retornou resultados.")
        r0 = yolo_results[0]
        names = r0.names if hasattr(r0, "names") else {}

        boxes = r0.boxes
        if boxes is None or boxes.xyxy is None or len(boxes) == 0:
            raise ValueError("Nenhum parafuso detectado.")

        xyxy = boxes.xyxy.cpu().numpy().astype(int)
        cls_ids = boxes.cls.cpu().numpy().astype(int) if boxes.cls is not None else np.zeros((xyxy.shape[0],), dtype=int)
        confs = boxes.conf.cpu().numpy() if boxes.conf is not None else np.ones((xyxy.shape[0],), dtype=float)

        detections_payload = []
        H, W = img_bgr.shape[:2]

        for i, (x1, y1, x2, y2) in enumerate(xyxy):
            # leve inset para evitar etiqueta e borda
            inset = 0.05
            w = x2 - x1
            h = y2 - y1
            x1 += int(w * inset)
            y1 += int(h * inset)
            x2 -= int(w * inset)
            y2 -= int(h * inset)

            x1 = max(0, min(x1, W - 1))
            y1 = max(0, min(y1, H - 1))
            x2 = max(x1 + 1, min(x2, W))
            y2 = max(y1 + 1, min(y2, H))

            crop_bgr = img_bgr[y1:y2, x1:x2].copy()
            if crop_bgr.size == 0:
                continue

            try:
                analysis, corrosion_rgb = process_image_bgr(crop_bgr, corrosion_type=corrosion_type)
            except Exception as e:
                analysis = {"error": f"Falha na análise do crop {i}: {e}", "corrosion_type": corrosion_type}
                corrosion_rgb = None

            s3_key_crop = None
            if corrosion_rgb is not None:
                bgr_result = cv2.cvtColor(corrosion_rgb, cv2.COLOR_RGB2BGR)
                ok, buffer = cv2.imencode('.png', bgr_result)
                if ok:
                    s3_key_crop = f"imagens_resultados/{uuid.uuid4()}.png"
                    s3_client.put_object(
                        Bucket=AWS_S3_BUCKET_NAME,
                        Key=s3_key_crop,
                        Body=buffer.tobytes(),
                        ContentType='image/png'
                    )
                    s3_keys_crops.append(s3_key_crop)

            cls_id = int(cls_ids[i]) if i < len(cls_ids) else 0
            label = names.get(cls_id, f"class_{cls_id}")
            score = float(confs[i]) if i < len(confs) else 0.0

            detections_payload.append({
                "index": i,
                "bbox_xyxy": [int(x1), int(y1), int(x2), int(y2)],
                "class_id": cls_id,
                "class_name": label or "Parafuso",
                "score": round(score, 4),
                "analysis": analysis,
                "s3_result_key": s3_key_crop
            })

        if not detections_payload:
            raise ValueError("Nenhum crop válido para análise.")

        # upload do original
        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}"
        s3_client.put_object(
            Bucket=AWS_S3_BUCKET_NAME,
            Key=s3_key_original,
            Body=content,
            ContentType=content_type
        )

        # imagem anotada
        labels_for_draw = [f"{d['class_name']} {d['score']:.2f}" for d in detections_payload]
        annotated_rgb = draw_boxes_on_bgr(img_bgr, xyxy, labels_for_draw)
        annotated_data_uri = _to_data_uri_from_rgb(annotated_rgb)

        # resumo para Supabase
        valid_percents = [d["analysis"].get("percent") for d in detections_payload if isinstance(d.get("analysis"), dict) and d["analysis"].get("percent") is not None]
        avg_percent = round(float(np.mean(valid_percents)), 4) if valid_percents else 0.0
        first_result_key = next((d["s3_result_key"] for d in detections_payload if d.get("s3_result_key")), None)

        dados_para_inserir = {
            "nome_amostra": file.filename,
            "percentual_corrosao": avg_percent,
            "pixels_totais_obj": None,
            "pixels_corrosao": None,
            "imagem_original": s3_key_original,
            "imagem_resultado": first_result_key
        }
        response = supabase_client.from_("amostras").insert(dados_para_inserir).execute()
        new_record_id = response.data[0]['id'] if response and response.data else None

        out = {
            "corrosion_type": corrosion_type,
            "database_id": new_record_id,
            "detections_count": len(detections_payload),
            "annotated_image": annotated_data_uri,
            "detections": detections_payload
        }
        return JSONResponse(content=out)

    except Exception as e:
        import traceback
        print("Erro no /analyze:", repr(e))
        traceback.print_exc()
        if s3_key_original:
            try:
                s3_client.delete_object(Bucket=AWS_S3_BUCKET_NAME, Key=s3_key_original)
            except Exception:
                pass
        for key in s3_keys_crops:
            try:
                s3_client.delete_object(Bucket=AWS_S3_BUCKET_NAME, Key=key)
            except Exception:
                pass
        raise HTTPException(status_code=500, detail=f"Erro interno: {e}")

@app.get("/samples/{sample_id}")
async def get_sample_images(sample_id: int):
    try:
        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 {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
            )
        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
            )
        return JSONResponse(content=links)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Erro ao buscar links: {e}")