|
|
import gradio as gr |
|
|
import pandas as pd |
|
|
import torch |
|
|
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration |
|
|
from PIL import Image |
|
|
import io |
|
|
from datetime import datetime |
|
|
|
|
|
|
|
|
model_name = "Qwen/Qwen2-VL-7B-Instruct" |
|
|
processor = AutoProcessor.from_pretrained(model_name) |
|
|
model = Qwen2VLForConditionalGeneration.from_pretrained( |
|
|
model_name, |
|
|
torch_dtype=torch.bfloat16, |
|
|
device_map="auto" |
|
|
) |
|
|
|
|
|
def extract_data_from_image(images): |
|
|
results = [] |
|
|
|
|
|
for idx, img_file in enumerate(images): |
|
|
try: |
|
|
image = Image.open(io.BytesIO(img_file.read())).convert("RGB") |
|
|
|
|
|
|
|
|
prompt = """ |
|
|
กรุณาสกัดข้อมูลสำคัญจากเอกสารนี้: |
|
|
- วันที่ |
|
|
- ยอดรวม |
|
|
- ชื่อร้านค้า |
|
|
- เลขใบเสร็จ |
|
|
|
|
|
กรุณาตอบในรูปแบบ JSON |
|
|
""" |
|
|
|
|
|
messages = [ |
|
|
{ |
|
|
"role": "user", |
|
|
"content": [ |
|
|
{"type": "image"}, |
|
|
{"type": "text", "text": prompt} |
|
|
] |
|
|
} |
|
|
] |
|
|
|
|
|
text_prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
|
|
inputs = processor(text=text_prompt, images=image, return_tensors="pt").to(model.device).bfloat16() |
|
|
|
|
|
with torch.no_grad(): |
|
|
generated_ids = model.generate(**inputs, max_new_tokens=512) |
|
|
|
|
|
generated_ids_trimmed = [out_ids[len(inputs["input_ids"][0]):] for out_ids in generated_ids] |
|
|
answer = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] |
|
|
|
|
|
try: |
|
|
structured = eval(answer.replace("```json", "").replace("```", "")) |
|
|
except: |
|
|
structured = {"raw_response": answer} |
|
|
|
|
|
results.append({ |
|
|
"file_name": img_file.name, |
|
|
"data": str(structured), |
|
|
"timestamp": datetime.now().isoformat() |
|
|
}) |
|
|
|
|
|
except Exception as e: |
|
|
results.append({ |
|
|
"file_name": img_file.name, |
|
|
"data": f"เกิดข้อผิดพลาด: {str(e)}", |
|
|
"timestamp": datetime.now().isoformat() |
|
|
}) |
|
|
|
|
|
df = pd.DataFrame(results) |
|
|
df["structured_data"] = df["data"].astype(str) |
|
|
|
|
|
|
|
|
parquet_path = "output.parquet" |
|
|
df.to_parquet(parquet_path) |
|
|
|
|
|
return { |
|
|
"table": df[["file_name", "structured_data"]], |
|
|
"download": parquet_path |
|
|
} |
|
|
|
|
|
|
|
|
title = "📄 ระบบสกัดข้อมูลเอกสารอัตโนมัติ (รองรับภาษาไทย)" |
|
|
description = "อัปโหลดภาพหลายไฟล์ → สกัดข้อมูล → แยกหัวข้อ → บันทึกเป็น Parquet" |
|
|
|
|
|
interface = gr.Interface( |
|
|
fn=extract_data_from_image, |
|
|
inputs=gr.File(type="file", file_types=["image"], multiple=True), |
|
|
outputs=[ |
|
|
gr.Dataframe(label="ผลลัพธ์"), |
|
|
gr.File(label="ดาวน์โหลด Parquet") |
|
|
], |
|
|
title=title, |
|
|
description=description, |
|
|
allow_flagging="never" |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
interface.launch() |
