protae5544 commited on
Commit
2a24fae
·
verified ·
1 Parent(s): dbf0579

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +102 -0
app.py ADDED
@@ -0,0 +1,102 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import pandas as pd
3
+ import torch
4
+ from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
5
+ from PIL import Image
6
+ import io
7
+ from datetime import datetime
8
+
9
+ # โหลดโมเดลและ processor จาก Hugging Face
10
+ model_name = "Qwen/Qwen2-VL-7B-Instruct"
11
+ processor = AutoProcessor.from_pretrained(model_name)
12
+ model = Qwen2VLForConditionalGeneration.from_pretrained(
13
+ model_name,
14
+ torch_dtype=torch.bfloat16,
15
+ device_map="auto"
16
+ )
17
+
18
+ def extract_data_from_image(images):
19
+ results = []
20
+
21
+ for idx, img_file in enumerate(images):
22
+ try:
23
+ image = Image.open(io.BytesIO(img_file.read())).convert("RGB")
24
+
25
+ # Prompt บอกโมเดลว่าให้ทำอะไร
26
+ prompt = """
27
+ กรุณาสกัดข้อมูลสำคัญจากเอกสารนี้:
28
+ - วันที่
29
+ - ยอดรวม
30
+ - ชื่อร้านค้า
31
+ - เลขใบเสร็จ
32
+
33
+ กรุณาตอบในรูปแบบ JSON
34
+ """
35
+
36
+ messages = [
37
+ {
38
+ "role": "user",
39
+ "content": [
40
+ {"type": "image"},
41
+ {"type": "text", "text": prompt}
42
+ ]
43
+ }
44
+ ]
45
+
46
+ text_prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
47
+ inputs = processor(text=text_prompt, images=image, return_tensors="pt").to(model.device).bfloat16()
48
+
49
+ with torch.no_grad():
50
+ generated_ids = model.generate(**inputs, max_new_tokens=512)
51
+
52
+ generated_ids_trimmed = [out_ids[len(inputs["input_ids"][0]):] for out_ids in generated_ids]
53
+ answer = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
54
+
55
+ try:
56
+ structured = eval(answer.replace("```json", "").replace("```", ""))
57
+ except:
58
+ structured = {"raw_response": answer}
59
+
60
+ results.append({
61
+ "file_name": img_file.name,
62
+ "data": str(structured),
63
+ "timestamp": datetime.now().isoformat()
64
+ })
65
+
66
+ except Exception as e:
67
+ results.append({
68
+ "file_name": img_file.name,
69
+ "data": f"เกิดข้อผิดพลาด: {str(e)}",
70
+ "timestamp": datetime.now().isoformat()
71
+ })
72
+
73
+ df = pd.DataFrame(results)
74
+ df["structured_data"] = df["data"].astype(str)
75
+
76
+ # บันทึกเป็น Parquet
77
+ parquet_path = "output.parquet"
78
+ df.to_parquet(parquet_path)
79
+
80
+ return {
81
+ "table": df[["file_name", "structured_data"]],
82
+ "download": parquet_path
83
+ }
84
+
85
+ # UI Components
86
+ title = "📄 ระบบสกัดข้อมูลเอกสารอัตโนมัติ (รองรับภาษาไทย)"
87
+ description = "อัปโหลดภาพหลายไฟล์ → สกัดข้อมูล → แยกหัวข้อ → บันทึกเป็น Parquet"
88
+
89
+ interface = gr.Interface(
90
+ fn=extract_data_from_image,
91
+ inputs=gr.File(type="file", file_types=["image"], multiple=True),
92
+ outputs=[
93
+ gr.Dataframe(label="ผลลัพธ์"),
94
+ gr.File(label="ดาวน์โหลด Parquet")
95
+ ],
96
+ title=title,
97
+ description=description,
98
+ allow_flagging="never"
99
+ )
100
+
101
+ if __name__ == "__main__":
102
+ interface.launch()