kawaiipeace's picture
update code
f28d3f0
import os
import re
from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from pdf2image import convert_from_bytes
from PIL import Image
from transformers import pipeline
import torch
import gradio as gr
# -------------------------
# Load Hugging Face Model
# -------------------------
ocr_pipeline = pipeline(
task="document-question-answering",
model="scb10x/typhoon-ocr-7b",
device=0 if torch.cuda.is_available() else -1
)
# -------------------------
# OCR Wrapper
# -------------------------
def ocr_document(image):
result = ocr_pipeline(image)
return result[0]["generated_text"] if result and isinstance(result, list) else str(result)
# -------------------------
# Regex Field Extractor
# -------------------------
def extract_fields_regex(text: str) -> dict:
text = re.sub(r"<.*?>", "", text) # remove tags
text = re.sub(r"\s{2,}", " ", text) # collapse spaces
text = re.sub(r"\n{2,}", "\n", text) # collapse newlines
patterns = {
"เลขที่ผู้เสียภาษี": r"(?:TAX\s*ID|เลขที่ผู้เสียภาษี)[\s:\-\.]*([\d]{10,13})",
"เลขที่ใบกำกับภาษี": r"(?:TAX\s*INV\.?|เลขที่ใบกำกับภาษี|ใบกำกับ)[\s:\-\.]*([\d]{8,20})",
"จำนวนเงิน": r"(?:AMOUNT\s*THB|จำนวนเงิน|รวมเงิน)[\s:\-\.]*([\d,]+\.\d{2})",
"ราคาต่อลิตร": r"(?:Baht\/Litr\.?|Bath\/Ltr\.?|ราคาต่อลิตร|ราคา\/ลิตร|ราคาน้ำมัน)[\s:\-\.]*([\d,]+\.\d{2})",
"ลิตร": r"(?:Ltr\.?|Ltrs?\.?|ลิตร)[\s:\-\.]*([\d,]+\.\d{3})",
"ภาษีมูลค่าเพิ่ม": r"(?:VAT|ภาษีมูลค่าเพิ่ม)[\s:\-\.]*([\d,]+\.\d{2})",
"ยอดรวม": r"(?:TOTAL\s*THB|ยอดรวม|รวมทั้งสิ้น|รวมเงินทั้งสิ้น)[\s:\-\.]*([\d,]+\.\d{2})",
"วันที่": r"(?:DATE|วันที่|ออกใบกำกับวันที่)[\s:\-\.]*([\d]{2}/[\d]{2}/[\d]{2,4})",
}
results = {}
for field, pattern in patterns.items():
match = re.search(pattern, text, re.IGNORECASE)
results[field] = match.group(1).strip() if match else None
return results
# -------------------------
# PDF Handling
# -------------------------
def pdf_to_image(file_bytes: bytes) -> Image.Image:
images = convert_from_bytes(file_bytes)
return images[0] # first page only
# -------------------------
# FastAPI App
# -------------------------
app = FastAPI()
# Optional: Allow all CORS (customize for security)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
@app.post("/api/ocr_receipt")
async def ocr_receipt(file: UploadFile = File(...)):
content = await file.read()
try:
# Convert PDF to image or use image directly
if file.filename.lower().endswith(".pdf"):
image = pdf_to_image(content)
else:
image = Image.open(file.file)
text = ocr_document(image)
fields = extract_fields_regex(text)
return {
"raw_ocr": text,
"extracted_fields": fields,
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# -------------------------
# Gradio UI
# -------------------------
def gradio_interface(image_path):
if isinstance(image_path, str) and image_path.lower().endswith(".pdf"):
with open(image_path, "rb") as f:
image = pdf_to_image(f.read())
else:
image = image_path
text = ocr_document(image)
extracted = extract_fields_regex(text)
return text, extracted
with gr.Blocks() as demo:
gr.Markdown("# 🧾 OCR ใบเสร็จ (Thai Receipt Scanner)")
with gr.Row():
img = gr.Image(type="filepath", label="อัปโหลดไฟล์ PDF หรือรูปภาพ")
out_text = gr.Textbox(label="ข้อความทั้งหมด", lines=12)
out_fields = gr.JSON(label="ฟิลด์ที่ดึงออกมา")
btn = gr.Button("ประมวลผล")
btn.click(fn=gradio_interface, inputs=img, outputs=[out_text, out_fields])
# For Hugging Face Spaces — no uvicorn needed
demo.launch()