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 pydantic import BaseModel

# --- KHỞI TẠO ỨNG DỤNG VÀ CÁC BIẾN TOÀN CỤC ---

app = FastAPI(
    title="Vietnamese Citizen ID OCR & Face Extraction API",
    description="Một microservice để trích xuất thông tin và cắt ảnh chân dung từ CCCD. Sử dụng Lazy Loading cho model.",
    version="1.3.0-lazyload-packaged"
)

# Khởi tạo các biến model toàn cục là None. Chúng sẽ được tải sau.
idcard_extractor = None
face_cascade = None
model_lock = threading.Lock()  # Lock để đảm bảo model chỉ được tải 1 lần trong môi trường đa luồng


# --- HÀM TẢI MODEL (LAZY LOADING) ---

def load_models():
    """
    Hàm này chỉ được gọi một lần duy nhất khi có request đầu tiên.
    Nó tải tất cả các model AI nặng vào bộ nhớ.
    """
    global idcard_extractor, face_cascade

    # Sử dụng lock để ngăn chặn nhiều request cùng lúc cố gắng tải model (race condition)
    with model_lock:
        # Kiểm tra lại một lần nữa bên trong lock, nếu một luồng khác đã tải xong thì bỏ qua.
        if idcard_extractor is None:
            print("--- LAZY LOADING MODELS (FIRST REQUEST) ---")
            try:
                # Import Extractor ngay tại đây, không import ở đầu file
                from core.extractor import Extractor

                # 1. Tải model OCR (sẽ đọc từ các file cục bộ trong thư mục /models)
                print("Loading OCR models...")
                idcard_extractor = Extractor()
                print("CCCD Text Extractor loaded successfully.")

                # 2. Tải model nhận diện khuôn mặt
                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.")
                face_cascade = cv2.CascadeClassifier(face_cascade_path)
                print("Face cascade classifier loaded successfully.")

            except Exception as e:
                print(f"FATAL: Error during model loading: {e}")
                # Đặt lại thành None để các request sau biết rằng model đã tải thất bại
                idcard_extractor = None
                face_cascade = None
            print("--- MODEL LOADING COMPLETE ---")


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

class ExtractionResponse(BaseModel):
    ID_number: Optional[str] = None
    Name: Optional[str] = None
    Date_of_birth: 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("/")
def read_root():
    return {"message": "Welcome to the CCCD Extraction API. POST to /extract/ to process an image."}


@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.
    load_models()

    # 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)