Spaces:
Sleeping
Sleeping
| # 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 | |
| app = FastAPI(title="Detector de Corrosão Branca") | |
| # PARA PROTOTIPO: permitir todas origens. Em produção restrinja ao domínio do frontend. | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| def process_image_bytes(img_bytes: bytes): | |
| # lê bytes em numpy + OpenCV | |
| 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"} | |
| 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) | |
| def to_data_uri(img_arr): | |
| # img_arr: RGB uint8 | |
| bgr = cv2.cvtColor(img_arr, 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}" | |
| isolated_b64 = to_data_uri(isolated_rgb) | |
| corrosion_b64 = to_data_uri(corrosion_vis) | |
| return { | |
| "percent": round(percent, 4), | |
| "total_pixels": total_pixels, | |
| "corrosion_pixels": corrosion_pixels, | |
| "isolated_image": isolated_b64, | |
| "corrosion_image": corrosion_b64, | |
| } | |
| 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) | |
| def read_root(): | |
| return {"status": "ok"} |