phgrouptechs/

scanocr-identity-vi

import cv2
import torch
from PIL import Image
from ultralytics import YOLO
from vietocr.tool.predictor import Predictor
from vietocr.tool.config import Cfg

# 1. Initialize VietOCR (For Vietnamese language)
config = Cfg.load_config_from_name('vgg_transformer')
config['device'] = 'cuda' if torch.cuda.is_available() else 'cpu'
ocr_predictor = Predictor(config)

# 2. Download the YOLO model (exported to ONNX)
yolo_model = YOLO("yolo_v11_best.onnx")

def scan_id_card(image_path):
    img = cv2.imread(image_path)
    detections = yolo_model(img)[0]
    
    results = {"type": "Unknown", "data": {}}
    box_list = []

    # Collect the discovered boxes
    for box in detections.boxes:
        cls = int(box.cls[0])
        label = detections.names[cls]
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        
        crop = img[y1:y2, x1:x2]
        box_list.append({"label": label, "y1": y1, "crop": crop})

    # Sort by Y-coordinate to avoid address inversion
    box_list = sorted(box_list, key=lambda x: x["y1"])

    for item in box_list:
        label = item["label"]
        crop_pil = Image.fromarray(cv2.cvtColor(item["crop"], cv2.COLOR_BGR2RGB))
        text = ocr_predictor.predict(crop_pil)
        
        if label in results["data"]:
            results["data"][label] += " " + text
        else:
            results["data"][label] = text

    return results

# Test run
print(scan_id_card("path_to_your_image.jpg"))

{
    "success": true,
    "result": {
        "type": "CCCD Gắn Chip",
        "data": {
            "Số ID": "026xxxxxxxxx",
            "Họ tên": "NGUYỄN QUỐC VIỆT",
            "Ngày sinh": "21/02/1989",
            "Quê quán": "Sơn Lôi, Bình Xuyên, Vĩnh Phúc",
            "Địa chỉ thường trú": "Ấp An Viễn, Bình An, Long Thành, Đồng Nai"
        }
    }
}