fabdRb/Flutter_app

1import os2import re3from io import BytesIO4from typing import Optional, List, Tuple5 6import cv27import numpy as np8from PIL import Image, ImageOps9 10from dotenv import load_dotenv11from fastapi import FastAPI, UploadFile, File, HTTPException12from fastapi.middleware.cors import CORSMiddleware13 14from ultralytics import YOLO15from paddleocr import PaddleOCR16 17load_dotenv()18 19YOLO_MODEL_PATH = os.getenv("YOLO_MODEL_PATH", "plate.pt")20OCR_LANG = os.getenv("OCR_LANG", "en")21OCR_USE_ANGLE_CLS = os.getenv("OCR_USE_ANGLE_CLS", "true").lower() == "true"22 23# Regex Brasil:24# - Antiga: ABC123425# - Mercosul: ABC1D23 (o 5º pode ser letra ou número dependendo do OCR)26BR_OLD = re.compile(r"^[A-Z]{3}[0-9]{4}$")27BR_MERCOSUL = re.compile(r"^[A-Z]{3}[0-9][A-Z0-9][0-9]{2}$")28 29app = FastAPI(title="ANPR - YOLO + PaddleOCR")30 31app.add_middleware(32    CORSMiddleware,33    allow_origins=["*"],34    allow_methods=["*"],35    allow_headers=["*"],36)37 38yolo_model: Optional[YOLO] = None39ocr_engine: Optional[PaddleOCR] = None40 41 42@app.on_event("startup")43def startup():44    global yolo_model, ocr_engine45 46    # YOLO47    if not os.path.exists(YOLO_MODEL_PATH):48        print("📂 Arquivos na raiz:", os.listdir("."))49        raise RuntimeError(f"❌ Modelo YOLO não encontrado: {YOLO_MODEL_PATH}")50 51    yolo_model = YOLO(YOLO_MODEL_PATH)52    print(f"✅ YOLO carregado: {YOLO_MODEL_PATH}")53 54    # PaddleOCR55    ocr_engine = PaddleOCR(56        use_angle_cls=OCR_USE_ANGLE_CLS,57        lang=OCR_LANG,58        show_log=False,59    )60    print(f"✅ PaddleOCR carregado | lang={OCR_LANG} | angle_cls={OCR_USE_ANGLE_CLS}")61 62 63def _read_image_bytes(file_bytes: bytes) -> np.ndarray:64    """65    ✅ Lê imagem aplicando EXIF transpose (câmera do celular)66    Retorna em BGR (OpenCV).67    """68    pil = Image.open(BytesIO(file_bytes))69    pil = ImageOps.exif_transpose(pil)  # ✅ corrige rotação EXIF70    rgb = np.array(pil.convert("RGB"))71    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)72    return bgr73 74 75def _pick_best_plate_bbox(yolo_result) -> Optional[Tuple[int, int, int, int, float]]:76    """77    Retorna bbox (x1,y1,x2,y2,conf) da melhor detecção.78    """79    if yolo_result is None or yolo_result.boxes is None:80        return None81 82    boxes = yolo_result.boxes83    if len(boxes) == 0:84        return None85 86    best = None87    best_score = -1.088 89    for b in boxes:90        conf = float(b.conf.item()) if b.conf is not None else 0.091        x1, y1, x2, y2 = b.xyxy[0].tolist()92        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)93 94        area = max(0, x2 - x1) * max(0, y2 - y1)95        score = conf * (1.0 + area / 200000.0)  # conf + leve bônus por área96 97        if score > best_score:98            best_score = score99            best = (x1, y1, x2, y2, conf)100 101    return best102 103 104def _crop_with_padding(img_bgr: np.ndarray, bbox: List[int], pad: int = 10) -> np.ndarray:105    h, w = img_bgr.shape[:2]106    x1, y1, x2, y2 = bbox107    x1 = max(0, x1 - pad)108    y1 = max(0, y1 - pad)109    x2 = min(w, x2 + pad)110    y2 = min(h, y2 + pad)111    return img_bgr[y1:y2, x1:x2].copy()112 113 114def _preprocess_plate_for_ocr(plate_bgr: np.ndarray) -> np.ndarray:115    """116    Pré-processamento leve pra melhorar OCR de placa:117    - cinza118    - resize119    - contraste120    """121    gray = cv2.cvtColor(plate_bgr, cv2.COLOR_BGR2GRAY)122 123    # aumenta tamanho se estiver pequeno124    h, w = gray.shape[:2]125    if w < 300:126        scale = 2.0127        gray = cv2.resize(gray, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_CUBIC)128 129    # melhora contraste (CLAHE)130    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))131    gray = clahe.apply(gray)132 133    return gray134 135 136def _clean_plate_text(text: str) -> str:137    text = text.upper()138    text = re.sub(r"[^A-Z0-9]", "", text)139    return text140 141 142def _best_plate_candidate(candidates: List[Tuple[str, float]]) -> Tuple[str, float]:143    if not candidates:144        return "", 0.0145 146    valid = []147    for t, s in candidates:148        if BR_OLD.match(t) or BR_MERCOSUL.match(t):149            valid.append((t, s))150 151    if valid:152        valid.sort(key=lambda x: x[1], reverse=True)153        return valid[0]154 155    candidates.sort(key=lambda x: x[1], reverse=True)156    return candidates[0]157 158 159def _run_paddleocr(plate_img_gray: np.ndarray) -> Tuple[str, float]:160    if ocr_engine is None:161        return "", 0.0162 163    result = ocr_engine.ocr(plate_img_gray, cls=True)164 165    candidates = []166    try:167        for line in result:168            for item in line:169                txt = item[1][0]170                score = float(item[1][1])171                cleaned = _clean_plate_text(txt)172                if cleaned:173                    candidates.append((cleaned, score))174    except Exception:175        pass176 177    candidates_sorted = sorted(candidates, key=lambda x: x[1], reverse=True)178 179    merged = []180    if len(candidates_sorted) >= 2:181        t1, s1 = candidates_sorted[0]182        t2, s2 = candidates_sorted[1]183        merged.append((t1 + t2, min(s1, s2)))184 185    all_candidates = candidates_sorted + merged186 187    best_text, best_conf = _best_plate_candidate(all_candidates)188    return best_text, float(best_conf)189 190 191def _bbox_to_norm(bbox: List[int], w: int, h: int) -> List[float]:192    """193    ✅ Converte bbox pixel -> bbox normalizada (0..1)194    """195    x1, y1, x2, y2 = bbox196    return [197        x1 / w,198        y1 / h,199        x2 / w,200        y2 / h,201    ]202 203 204@app.get("/")205def root():206    return {"status": "ok", "message": "ANPR API online", "docs": "/docs"}207 208 209@app.post("/predict")210async def predict(file: UploadFile = File(...)):211    """212    Endpoint pro Flutter: envia 'file' multipart.213    Retorna bbox em pixels + bbox_norm (0..1)214    """215    if yolo_model is None:216        raise HTTPException(status_code=500, detail="YOLO não carregou.")217    if ocr_engine is None:218        raise HTTPException(status_code=500, detail="PaddleOCR não carregou.")219 220    file_bytes = await file.read()221    if not file_bytes:222        raise HTTPException(status_code=400, detail="Arquivo vazio.")223 224    try:225        img_bgr = _read_image_bytes(file_bytes)226    except Exception as e:227        raise HTTPException(status_code=400, detail=f"Erro ao ler imagem: {e}")228 229    h, w = img_bgr.shape[:2]230 231    # YOLO detect232    yres = yolo_model.predict(img_bgr, conf=0.25, verbose=False)233    yres0 = yres[0] if len(yres) > 0 else None234 235    best = _pick_best_plate_bbox(yres0)236    if best is None:237        return {238            "plate": "",239            "confidence": 0.0,240            "plate_model_conf": 0.0,241            "bbox": None,242            "bbox_norm": None,   # ✅ NOVO243            "image_w": w,        # ✅ opcional244            "image_h": h,        # ✅ opcional245            "view_used": "paddleocr",246        }247 248    x1, y1, x2, y2, det_conf = best249    bbox = [int(x1), int(y1), int(x2), int(y2)]250 251    # ✅ bbox normalizada (0..1)252    bbox_norm = _bbox_to_norm(bbox, w=w, h=h)253 254    # Crop placa (só pra OCR)255    crop = _crop_with_padding(img_bgr, bbox, pad=12)256 257    # Preprocess OCR258    crop_gray = _preprocess_plate_for_ocr(crop)259 260    # OCR261    plate_text, ocr_conf = _run_paddleocr(crop_gray)262 263    return {264        "plate": plate_text,265        "confidence": float(ocr_conf),266        "plate_model_conf": float(det_conf),267        "bbox": bbox,268        "bbox_norm": bbox_norm,  # ✅ NOVO269        "image_w": w,            # ✅ opcional270        "image_h": h,            # ✅ opcional271        "view_used": "paddleocr",272    }273