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