Update app.py

app.py

@@ -1,41 +1,259 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import base64
+import json
+import ast
+import re
+from io import BytesIO
+
 import torch
+import torchvision.transforms as T
+from PIL import Image
+from torchvision.transforms.functional import InterpolationMode
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoModel, AutoTokenizer
+
+
+app = FastAPI(title="CCCD OCR with Vintern-1B-v2")
 
-app = FastAPI()
 
 MODEL_NAME = "5CD-AI/Vintern-1B-v2"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
 
-print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-print("Loading model (INT4, CPU)...")
-model = AutoModelForCausalLM.from_pretrained(
+print(f"Loading model `{MODEL_NAME}` on {DEVICE} ...")
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_NAME,
+    trust_remote_code=True,
+    use_fast=False,
+)
+model = AutoModel.from_pretrained(
     MODEL_NAME,
-    load_in_4bit=True,
-    device_map="cpu",
-    torch_dtype=torch.float16
+    torch_dtype=DTYPE,
+    low_cpu_mem_usage=True,
+    trust_remote_code=True,
+)
+model.eval().to(DEVICE)
+
+generation_config = dict(
+    max_new_tokens=512,
+    do_sample=False,
+    num_beams=3,
+    repetition_penalty=3.5,
 )
 
-class InferRequest(BaseModel):
-    text: str
 
-@app.post("/infer")
-def infer(req: InferRequest):
-    inputs = tokenizer(
-        req.text,
-        return_tensors="pt",
-        truncation=True,
-        max_length=512
+# =========================
+# Image preprocessing (from notebook)
+# =========================
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+
+def build_transform(input_size: int):
+    mean, std = IMAGENET_MEAN, IMAGENET_STD
+    transform = T.Compose(
+        [
+            T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
+            T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
+            T.ToTensor(),
+            T.Normalize(mean=mean, std=std),
+        ]
     )
+    return transform
+
+
+def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+    best_ratio_diff = float("inf")
+    best_ratio = (1, 1)
+    area = width * height
+    for ratio in target_ratios:
+        target_aspect_ratio = ratio[0] / ratio[1]
+        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+        if ratio_diff < best_ratio_diff:
+            best_ratio_diff = ratio_diff
+            best_ratio = ratio
+        elif ratio_diff == best_ratio_diff:
+            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                best_ratio = ratio
+    return best_ratio
+
+
+def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
+    orig_width, orig_height = image.size
+    aspect_ratio = orig_width / orig_height
+
+    target_ratios = set(
+        (i, j)
+        for n in range(min_num, max_num + 1)
+        for i in range(1, n + 1)
+        for j in range(1, n + 1)
+        if i * j <= max_num and i * j >= min_num
+    )
+    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+    target_aspect_ratio = find_closest_aspect_ratio(
+        aspect_ratio, target_ratios, orig_width, orig_height, image_size
+    )
+
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    resized_img = image.resize((target_width, target_height))
+    processed_images = []
+    for i in range(blocks):
+        box = (
+            (i % (target_width // image_size)) * image_size,
+            (i // (target_width // image_size)) * image_size,
+            ((i % (target_width // image_size)) + 1) * image_size,
+            ((i // (target_width // image_size)) + 1) * image_size,
+        )
+        split_img = resized_img.crop(box)
+        processed_images.append(split_img)
+    assert len(processed_images) == blocks
+    if use_thumbnail and len(processed_images) != 1:
+        thumbnail_img = image.resize((image_size, image_size))
+        processed_images.append(thumbnail_img)
+    return processed_images
+
 
+def load_image_from_base64(base64_string: str, input_size=448, max_num=12):
+    if base64_string.startswith("data:image"):
+        base64_string = base64_string.split(",", 1)[1]
+
+    image_data = base64.b64decode(base64_string)
+    image = Image.open(BytesIO(image_data)).convert("RGB")
+    transform = build_transform(input_size=input_size)
+    images = dynamic_preprocess(
+        image, image_size=input_size, use_thumbnail=True, max_num=max_num
+    )
+    pixel_values = [transform(img) for img in images]
+    pixel_values = torch.stack(pixel_values)
+    return pixel_values
+
+
+# =========================
+# Prompt & helpers
+# =========================
+PROMPT = """<image>
+Bạn là hệ thống OCR + trích xuất dữ liệu từ ảnh Căn cước công dân (CCCD) Việt Nam.
+Nhiệm vụ: đọc đúng chữ trên thẻ và trả về CHỈ 1 đối tượng JSON theo schema quy định.
+
+QUY TẮC BẮT BUỘC:
+1) Chỉ trả về JSON thuần (không markdown, không giải thích, không thêm ký tự nào ngoài JSON).
+2) Chỉ được có đúng 5 khóa sau (đúng chính tả, đúng chữ thường, có dấu gạch dưới):
+   - "so_no"
+   - "ho_va_ten"
+   - "ngay_sinh"
+   - "que_quan"
+   - "noi_thuong_tru"
+   Không được thêm bất kỳ khóa nào khác.
+3) Mapping trường (lấy theo NHÃN in trên thẻ, không lấy từ QR):
+   - so_no: lấy giá trị ngay sau nhãn "Số / No." (hoặc "Số/No.").
+   - ho_va_ten: lấy giá trị ngay sau nhãn "Họ và tên / Full name".
+   - ngay_sinh: lấy giá trị ngay sau nhãn "Ngày sinh / Date of birth"; nếu có định dạng dd/mm/yyyy thì giữ đúng dd/mm/yyyy.
+   - que_quan: lấy giá trị ngay sau nhãn "Quê quán / Place of origin".
+   - noi_thuong_tru: lấy giá trị ngay sau nhãn "Nơi thường trú / Place of residence".
+4) Nếu trường nào không đọc được rõ/chắc chắn: đặt null. Không được suy đoán.
+5) Chuẩn hoá: trim khoảng trắng đầu/cuối; giữ nguyên dấu tiếng Việt và chữ hoa/thường như trong ảnh.
+
+CHỈ TRẢ VỀ THEO MẪU JSON NÀY:
+{
+  "so_no": "... hoặc null",
+  "ho_va_ten": "... hoặc null",
+  "ngay_sinh": "... hoặc null",
+  "que_quan": "... hoặc null",
+  "noi_thuong_tru": "... hoặc null"
+}
+"""
+
+
+def parse_response_to_json(response_text: str):
+    if not response_text:
+        return None
+
+    s = response_text.strip()
+
+    if s.startswith('"') and s.endswith('"'):
+        s = s[1:-1].replace('\\"', '"')
+
+    try:
+        obj = json.loads(s)
+        if isinstance(obj, dict):
+            return obj
+    except json.JSONDecodeError:
+        pass
+
+    try:
+        obj = ast.literal_eval(s)
+        if isinstance(obj, dict):
+            return obj
+    except (ValueError, SyntaxError):
+        pass
+
+    json_pattern = r"\{[\s\S]*\}"
+    m = re.search(json_pattern, s)
+    if m:
+        chunk = m.group(0).strip()
+        try:
+            obj = ast.literal_eval(chunk)
+            if isinstance(obj, dict):
+                return obj
+        except Exception:
+            pass
+        try:
+            chunk2 = chunk.replace("'", '"')
+            obj = json.loads(chunk2)
+            if isinstance(obj, dict):
+                return obj
+        except Exception:
+            pass
+
+    return {"text": response_text}
+
+
+def normalize_base64(image_base64: str) -> str:
+    if not image_base64:
+        return image_base64
+    image_base64 = image_base64.strip()
+    if image_base64.startswith("data:"):
+        parts = image_base64.split(",", 1)
+        if len(parts) == 2:
+            return parts[1]
+    return image_base64
+
+
+def ocr_by_llm(image_base64: str, prompt: str) -> str:
+    pixel_values = load_image_from_base64(image_base64, max_num=6)
+    if DEVICE == "cuda":
+        pixel_values = pixel_values.to(dtype=torch.bfloat16, device=DEVICE)
+    else:
+        pixel_values = pixel_values.to(dtype=torch.float32, device=DEVICE)
     with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            do_sample=False
+        response_message = model.chat(
+            tokenizer,
+            pixel_values,
+            prompt,
+            generation_config,
         )
+    del pixel_values
+    return response_message
+
+
+class OCRRequest(BaseModel):
+    image_base64: str
+
+
+@app.post("/ocr")
+def ocr_endpoint(req: OCRRequest):
+    image_base64 = normalize_base64(req.image_base64)
+    if not image_base64:
+        raise HTTPException(status_code=400, detail="image_base64 is required")
 
-    result = tokenizer.decode(output[0], skip_special_tokens=True)
-    return {"result": result}
+    try:
+        response_message = ocr_by_llm(image_base64, PROMPT)
+        parsed = parse_response_to_json(response_message)
+        return {"response_message": parsed}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))