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