app.py

import base64
import datetime
import os
import re
import time
import uuid

import cv2
from fastapi import FastAPI, File, UploadFile, HTTPException
from pydantic import BaseModel
from typing import Optional

# Import lớp Extractor từ thư mục core
from core.extractor import Extractor

# --- Khởi tạo ---

# Khởi tạo ứng dụng FastAPI
app = FastAPI(
    title="CCCD Extraction API",
    description="Một microservice để trích xuất thông tin từ Căn cước công dân Việt Nam.",
    version="1.0.0"
)

# Đường dẫn để lưu trữ file upload
UPLOAD_DIR = "uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)

# Khởi tạo một lần duy nhất đối tượng Extractor để tái sử dụng
# Điều này giúp load model một lần và tăng tốc độ xử lý cho các request sau
try:
    idcard_extractor = Extractor()
    print("CCCD Extractor loaded successfully.")
except Exception as e:
    print(f"Error loading CCCD Extractor: {e}")
    idcard_extractor = None


# --- Định nghĩa Model cho Request và Response ---

# Model cho request nếu gửi ảnh dạng base64
class ImageRequest(BaseModel):
    image_base64: str


# Model cho response trả về
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
    elapsed: float


# --- Xây dựng API Endpoint ---

@app.get("/")
def read_root():
    return {"message": "Welcome to the CCCD Extraction API. Use the /extract/ endpoint to process an image."}


@app.post("/extract/", response_model=ExtractionResponse, tags=["CCCD Extraction"])
async def extract_id_card_info(file: UploadFile = File(...)):
    """
    Nhận một file ảnh CCCD, trích xuất thông tin và trả về.
    """
    if not idcard_extractor:
        raise HTTPException(status_code=500, detail="OCR Extractor is not available.")

    # --- 1. Lưu file ảnh được upload ---
    # Tạo tên file ngẫu nhiên và an toàn để tránh trùng lặp
    file_extension = os.path.splitext(file.filename)[1]
    random_filename = f"{uuid.uuid4()}{file_extension}"
    file_path = os.path.join(UPLOAD_DIR, random_filename)

    try:
        # Đọc nội dung file và lưu lại
        with open(file_path, "wb") as buffer:
            buffer.write(await file.read())
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Could not save uploaded file: {e}")

    # --- 2. Xử lý ảnh và trích xuất thông tin (logic từ Django view) ---
    start_time = time.time()
    try:
        frame = cv2.imread(file_path)
        if frame is None:
            raise HTTPException(status_code=400, detail="Invalid image file.")

        # Bước 1: Dùng PaddleOCR để phát hiện các vùng văn bản
        annotations = idcard_extractor.Detection(frame)

        info = {}

        # Tìm số CCCD trước tiên
        for box in annotations:
            text_detected = box[1][0]
            if re.search(r'\d{9,12}', text_detected):
                # Tách số ra khỏi chuỗi nhiễu
                id_number = re.search(r'\d{9,12}', text_detected).group(0)
                info['ID_number'] = id_number
                info['ID_number_box'] = box[0]
                break

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

        # Bước 2: Dùng VietOCR để nhận dạng các trường thông tin còn lại
        extracted_result = []
        for box in annotations:
            # Bỏ qua vùng chứa số ID đã xử lý
            if re.search(r'\d{9,12}', box[1][0]):
                continue

            top_left = (int(box[0][0][0]), int(box[0][0][1]))
            top_right = (int(box[0][1][0]), int(box[0][1][1]))
            bottom_right = (int(box[0][2][0]), int(box[0][2][1]))
            bottom_left = (int(box[0][3][0]), int(box[0][3][1]))

            # Warp và nhận dạng
            result_text, _ = idcard_extractor.WarpAndRec(frame, top_left, top_right, bottom_right, bottom_left)
            extracted_result.append((result_text, box[0]))  # Lưu cả text và bounding box

        # Bước 3: Tổng hợp thông tin
        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)

        # Xóa file tạm sau khi xử lý xong
        os.remove(file_path)

        return final_info

    except Exception as e:
        # Nếu có lỗi, cũng xóa file tạm
        if os.path.exists(file_path):
            os.remove(file_path)
        raise HTTPException(status_code=500, detail=f"An error occurred during processing: {str(e)}")