app.py

# app.py

import base64
import os
import re
import time
import uuid
import threading
from typing import Optional

import cv2
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import RedirectResponse
from pydantic import BaseModel
from core.extractor import Extractor

# --- KHỞI TẠO ỨNG DỤNG VÀ CÁC BIẾN TOÀN CỤC ---
description_md = """
### Microservice trích xuất thông tin CCCD 🔎

API này sử dụng các thư viện `vietocr` và `paddleocr` để thực hiện các chức năng sau:

1.  **Nhận diện và đọc (OCR)** các trường thông tin trên ảnh Căn cước công dân.
2.  **Phát hiện và cắt ảnh chân dung** từ ảnh CCCD gốc.
3.  Sử dụng cơ chế **Lazy Loading** để tối ưu hóa thời gian khởi động server.

_API được xây dựng với FastAPI._
"""

ml_models = {}

app = FastAPI(
    title="Vietnamese Citizen ID OCR API",
    description=description_md,
    version="1.4.0",
)

@app.on_event("startup")
async def startup_event():
    """
    Code này sẽ chạy KHI server khởi động.
    Thực hiện tải các model AI (Eager Loading).
    """
    print("--- EAGER LOADING MODELS (ON STARTUP) ---")
    try:
        print("Loading OCR models...")
        ml_models["idcard_extractor"] = Extractor()
        print("CCCD Text Extractor loaded successfully.")

        print("Loading face detection model...")
        face_cascade_path = os.path.join(cv2.data.haarcascades, 'haarcascade_frontalface_default.xml')
        if not os.path.exists(face_cascade_path):
            raise FileNotFoundError("Không tìm thấy file haarcascade.")
        ml_models["face_cascade"] = cv2.CascadeClassifier(face_cascade_path)
        print("Face cascade classifier loaded successfully.")

        print("--- MODEL LOADING COMPLETE ---")
    except Exception as e:
        print(f"FATAL: Error during model loading on startup: {e}")
        # Nếu có lỗi, re-raise exception để ngăn server khởi động
        raise

@app.on_event("shutdown")
async def shutdown_event():
    """
    Code này sẽ chạy KHI server tắt (shutdown).
    Dọn dẹp các model.
    """
    print("--- Cleaning up models ---")
    ml_models.clear()

# --- ĐỊNH NGHĨA MODEL CHO RESPONSE ---

class ExtractionResponse(BaseModel):
    ID_number: Optional[str] = None
    Name: Optional[str] = None
    Date_of_birth: Optional[str] = None
    Date_of_issue: Optional[str] = None
    Gender: Optional[str] = None
    Nationality: Optional[str] = None
    Place_of_origin: Optional[str] = None
    Place_of_residence: Optional[str] = None
    portrait_image_base64: Optional[str] = None
    elapsed: float

# --- API ENDPOINT ---
@app.get("/", include_in_schema=False)
async def root():
    """
    Khi người dùng truy cập trang gốc, tự động chuyển hướng đến trang tài liệu API.
    """
    return RedirectResponse(url="/docs")


@app.post("/extract/", response_model=ExtractionResponse, tags=["CCCD Extraction"])
async def extract_id_card_info(file: UploadFile = File(...)):
    """
    Nhận ảnh CCCD, trích xuất thông tin và cắt ảnh chân dung.
    Tải các model AI nếu đây là request đầu tiên.
    """
    # Bước 1: Tải model nếu chưa có
    # Nếu model đã được tải, hàm này sẽ bỏ qua rất nhanh.
    idcard_extractor = ml_models.get("idcard_extractor")
    face_cascade = ml_models.get("face_cascade")

    # Kiểm tra xem model đã được tải thành công chưa
    if not idcard_extractor or not face_cascade:
        raise HTTPException(status_code=503,
                            detail="Server is starting or models failed to load. Please try again in a moment.")

    # Bước 2: Tạo thư mục upload tạm thời trong /tmp và xác định đường dẫn file
    upload_dir = "/tmp/uploads"
    os.makedirs(upload_dir, exist_ok=True)
    file_path = os.path.join(upload_dir, f"{uuid.uuid4()}{os.path.splitext(file.filename)[1]}")

    start_time = time.time()
    try:
        # Bước 3: Lưu file ảnh được upload
        with open(file_path, "wb") as buffer:
            buffer.write(await file.read())

        frame = cv2.imread(file_path)
        if frame is None:
            raise HTTPException(status_code=400, detail="Invalid image file.")

        # Bước 4: Nhận diện và cắt ảnh chân dung
        gray_image = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        gray_image = cv2.equalizeHist(gray_image)
        faces = face_cascade.detectMultiScale(gray_image, scaleFactor=1.1, minNeighbors=5, minSize=(80, 80))

        portrait_base64 = None
        if len(faces) > 0:
            faces = sorted(faces, key=lambda f: f[2] * f[3], reverse=True)
            (x, y, w, h) = faces[0]
            padding_y, padding_x = int(h * 0.2), int(w * 0.2)
            portrait_img = frame[max(0, y - padding_y):min(frame.shape[0], y + h + padding_y),
                           max(0, x - padding_x):min(frame.shape[1], x + w + padding_x)]
            _, buffer = cv2.imencode('.jpg', portrait_img)
            portrait_base64 = base64.b64encode(buffer).decode('utf-8')

        # Bước 5: Trích xuất thông tin văn bản
        annotations = idcard_extractor.Detection(frame)
        info = {}
        for box in annotations:
            text_detected = box[1][0]
            id_match = re.search(r'\d{9,12}', text_detected)
            if id_match:
                info['ID_number'] = id_match.group(0)
                info['ID_number_box'] = box[0]
                break

        if 'ID_number' not in info:
            raise HTTPException(status_code=400, detail="Could not detect ID number.")

        extracted_result = []
        for box in annotations:
            if re.search(r'\d{9,12}', box[1][0]): continue
            top_left, top_right, bottom_right, bottom_left = (
            tuple(map(int, box[0][0])), tuple(map(int, box[0][1])), tuple(map(int, box[0][2])),
            tuple(map(int, box[0][3])))
            result_text, _ = idcard_extractor.WarpAndRec(frame, top_left, top_right, bottom_right, bottom_left)
            extracted_result.append((result_text, box[0]))

        # Bước 6: Tổng hợp kết quả và trả về
        final_info = idcard_extractor.GetInformationAndSave(extracted_result, info['ID_number'], info['ID_number_box'])
        elapsed = time.time() - start_time
        final_info["elapsed"] = round(elapsed, 2)
        final_info["portrait_image_base64"] = portrait_base64

        return final_info

    except Exception as e:
        # Ghi lại lỗi chi tiết vào log của server để gỡ lỗi
        print(f"Error during extraction: {e}")
        raise HTTPException(status_code=500, detail=f"An error occurred during processing: {str(e)}")
    finally:
        # Bước 7: Dọn dẹp file tạm sau khi xử lý xong
        if os.path.exists(file_path):
            os.remove(file_path)