Create app3.py

app3.py

@@ -0,0 +1,430 @@
+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}")