112 lines
4.6 KiB
Python
112 lines
4.6 KiB
Python
import io
|
|
import logging
|
|
import base64
|
|
import httpx
|
|
import fitz
|
|
from PIL import Image
|
|
from typing import List, Tuple
|
|
from core.models import OCRPageResult
|
|
from core.config import settings
|
|
|
|
logger = logging.getLogger("OCRService")
|
|
|
|
class OCRService:
|
|
"""
|
|
OCR Service implementation acting as a VLM client.
|
|
"""
|
|
def __init__(self):
|
|
self.vlm_url = settings.VLM_ENDPOINT
|
|
logger.info(f"Initialized VLM OCR Service connecting to {self.vlm_url}")
|
|
|
|
def _image_to_base64(self, img: Image.Image) -> str:
|
|
"""Chuyển đổi PIL Image sang chuẩn Base64 JPEG"""
|
|
buffered = io.BytesIO()
|
|
# Chuyển sang RGB nếu ảnh có kênh Alpha
|
|
if img.mode != 'RGB':
|
|
img = img.convert('RGB')
|
|
img.save(buffered, format="JPEG", quality=85)
|
|
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
|
return f"data:image/jpeg;base64,{img_str}"
|
|
|
|
def process_pdf_bytes(self, pdf_bytes: bytes) -> List[OCRPageResult]:
|
|
"""Process a PDF from memory using Vintern-3B VLM via LAN"""
|
|
if not pdf_bytes:
|
|
logger.warning("Empty PDF bytes received.")
|
|
return []
|
|
|
|
results = []
|
|
try:
|
|
import gc
|
|
doc = fitz.open(stream=pdf_bytes, filetype="pdf")
|
|
for page_num in range(len(doc)):
|
|
logger.info(f"VLM Processing page {page_num + 1}/{len(doc)} via LAN...")
|
|
|
|
# Render trang PDF thành ảnh. Hạ độ phân giải xuống 1.2 để giảm thiểu số lượng token
|
|
# Tránh lỗi 500 do vượt quá Context Window của Llama.cpp
|
|
matrix = fitz.Matrix(1.2, 1.2)
|
|
pix = doc[page_num].get_pixmap(matrix=matrix)
|
|
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
|
|
|
|
del pix
|
|
gc.collect()
|
|
|
|
# Chuyển ảnh sang Base64
|
|
b64_image = self._image_to_base64(img)
|
|
|
|
# Gọi API Llama.cpp Server
|
|
payload = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": b64_image
|
|
}
|
|
},
|
|
{
|
|
"type": "text",
|
|
"text": "Hãy trích xuất chính xác toàn bộ văn bản có trong hình ảnh này. Giữ nguyên định dạng và các dấu câu tiếng Việt."
|
|
}
|
|
]
|
|
}
|
|
],
|
|
"temperature": settings.VLM_TEMPERATURE,
|
|
"max_tokens": settings.VLM_MAX_TOKENS
|
|
}
|
|
|
|
try:
|
|
with httpx.Client(timeout=settings.VLM_TIMEOUT) as client:
|
|
response = client.post(self.vlm_url, json=payload)
|
|
response.raise_for_status()
|
|
|
|
data = response.json()
|
|
vlm_text = data['choices'][0]['message']['content'].strip()
|
|
|
|
results.append(OCRPageResult(
|
|
page=page_num + 1,
|
|
text=vlm_text,
|
|
confidence=0.99, # VLM thường không trả về độ tự tin từng chữ, set cứng 0.99
|
|
paddle_text="", # Bỏ qua cột so sánh cũ
|
|
paddle_confidence=0.0
|
|
))
|
|
logger.info(f"VLM extraction successful for page {page_num + 1}")
|
|
|
|
except Exception as api_err:
|
|
logger.error(f"VLM API Error: {api_err}")
|
|
# Ghi nhận trang lỗi nhưng vẫn tiếp tục các trang sau
|
|
results.append(OCRPageResult(
|
|
page=page_num + 1,
|
|
text=f"[LỖI KẾT NỐI VLM: {api_err}]",
|
|
confidence=0.0,
|
|
paddle_text="",
|
|
paddle_confidence=0.0
|
|
))
|
|
|
|
return results
|
|
except Exception as e:
|
|
import traceback
|
|
logger.error(f"Failed to process PDF: {e}\n{traceback.format_exc()}")
|
|
return []
|