ranbac commited on
Commit
f082006
·
verified ·
1 Parent(s): 7dd9372

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +44 -60
app.py CHANGED
@@ -17,50 +17,21 @@ import requests
17
  # Tắt log thừa
18
  logging.getLogger("ppocr").setLevel(logging.WARNING)
19
 
20
- # --- QUẢN LÝ MODEL ĐA NGÔN NGỮ ---
21
  print("Đang khởi tạo hệ thống quản lý Model...")
22
 
23
- # 1. Định nghĩa Mapping từ tên ngôn ngữ sang PaddleOCR
24
- # Danh sách này bao gồm các ngôn ngữ bạn yêu cầu
25
- LANG_MAP = {
26
- "Vietnamese (Tiếng Việt)": "vi",
27
- "English (Tiếng Anh)": "en",
28
- "Chinese (Trung Quốc)": "ch",
29
- "French (Pháp)": "fr",
30
- "German (Đức)": "de",
31
- "Korean (Hàn Quốc)": "korean",
32
- "Japan (Nhật Bản)": "japan",
33
- "Spanish (Tây Ban Nha)": "es",
34
- "Portuguese (Bồ Đào Nha)": "pt",
35
- "Italian (Ý)": "it",
36
- "Russian (Nga)": "ru",
37
- # Các ngôn ngữ Latin mở rộng (Mapping về mã code cụ thể hoặc 'latin' nếu model chung)
38
- "Afrikaans": "af", "Albanian": "sq", "Azerbaijani": "az", "Basque": "eu", "Bosnian": "bs",
39
- "Catalan": "ca", "Croatian": "hr", "Czech": "cs", "Danish": "da", "Dutch": "nl",
40
- "Estonian": "et", "Finnish": "fi", "Galician": "gl", "Hungarian": "hu", "Icelandic": "is",
41
- "Indonesian": "id", "Irish": "ga", "Kurdish": "ku", "Latin": "la", "Latvian": "lv",
42
- "Lithuanian": "lt", "Malay": "ms", "Maltese": "mt", "Maori": "mi", "Norwegian": "no",
43
- "Occitan": "oc", "Polish": "pl", "Romanian": "ro", "Slovak": "sk", "Slovenian": "sl",
44
- "Swahili": "sw", "Swedish": "sv", "Tagalog": "tl", "Turkish": "tr", "Uzbek": "uz",
45
- "Welsh": "cy",
46
- # Fallback cho các ngôn ngữ hiếm khác vào model Latin chung
47
- "General Latin (Other)": "latin"
48
- }
49
-
50
- # Cache để lưu các model đã tải
51
  OCR_ENGINES = {}
52
 
53
- def get_ocr_model(lang_name):
54
- # Lấy code từ tên hiển thị, mặc định 'en' nếu không tìm thấy
55
- lang_code = LANG_MAP.get(lang_name, 'en')
56
-
57
- # Nếu model này đã load rồi thì dùng lại ngay
58
  if lang_code in OCR_ENGINES:
59
  return OCR_ENGINES[lang_code]
60
 
61
- print(f"🔄 Đang tải model ngôn ngữ: {lang_code}...")
62
  try:
63
  # Khởi tạo PaddleOCR với ngôn ngữ được chọn
 
64
  engine = PaddleOCR(
65
  use_textline_orientation=True,
66
  use_doc_orientation_classify=False,
@@ -70,8 +41,8 @@ def get_ocr_model(lang_name):
70
  OCR_ENGINES[lang_code] = engine
71
  return engine
72
  except Exception as e:
73
- print(f"⚠️ Lỗi khởi tạo ngôn ngữ {lang_code}: {e}. Chuyển về chế độ Latin mặc định.")
74
- # Fallback an toàn
75
  engine = PaddleOCR(lang='en')
76
  OCR_ENGINES[lang_code] = engine
77
  return engine
@@ -95,9 +66,11 @@ FONT_PATH = check_and_download_font()
95
  def universal_draw(image, raw_data, font_path):
96
  if image is None: return image
97
 
 
98
  if isinstance(image, np.ndarray):
99
  image = Image.fromarray(image)
100
 
 
101
  canvas = image.copy()
102
  draw = ImageDraw.Draw(canvas)
103
 
@@ -107,6 +80,7 @@ def universal_draw(image, raw_data, font_path):
107
  except:
108
  font = ImageFont.load_default()
109
 
 
110
  def parse_box(b):
111
  try:
112
  if hasattr(b, 'tolist'): b = b.tolist()
@@ -118,6 +92,8 @@ def universal_draw(image, raw_data, font_path):
118
 
119
  items_to_draw = []
120
 
 
 
121
  processed = False
122
  if isinstance(raw_data, list) and len(raw_data) > 0 and isinstance(raw_data[0], dict):
123
  data_dict = raw_data[0]
@@ -131,6 +107,7 @@ def universal_draw(image, raw_data, font_path):
131
  if box and txt: items_to_draw.append((box, txt))
132
  processed = True
133
 
 
134
  if not processed:
135
  def hunt(data):
136
  if isinstance(data, dict):
@@ -150,9 +127,12 @@ def universal_draw(image, raw_data, font_path):
150
  for item in data: hunt(item)
151
  hunt(raw_data)
152
 
 
153
  for box, txt in items_to_draw:
154
  try:
 
155
  draw.polygon(box, outline="red", width=3)
 
156
  txt_x, txt_y = box[0]
157
  if hasattr(draw, "textbbox"):
158
  text_bbox = draw.textbbox((txt_x, txt_y), txt, font=font, anchor="lb")
@@ -189,43 +169,44 @@ def clean_text_result(text_list):
189
  cleaned.append(t)
190
  return cleaned
191
 
192
- # --- MAIN PREDICT (CẬP NHẬT: Nhận thêm tham số ngôn ngữ) ---
193
- def predict(image, selected_lang):
194
  if image is None: return None, "Chưa có ảnh.", "No Data"
195
 
196
  try:
197
- # 1. Lấy model dựa trên ngôn ngữ người dùng chọn
198
- current_ocr = get_ocr_model(selected_lang)
199
 
200
- # Chuẩn bị ảnh
201
  original_pil = image.copy() if isinstance(image, Image.Image) else Image.fromarray(image).copy()
202
  image_np = np.array(image)
203
 
204
- # 2. OCR
205
  raw_result = current_ocr.ocr(image_np)
206
 
207
- # 3. XỬ LÝ ẢNH ĐỂ VẼ (Giữ nguyên logic cũ)
208
  target_image_for_drawing = original_pil
 
 
209
  if isinstance(raw_result, list) and len(raw_result) > 0 and isinstance(raw_result[0], dict):
210
  if 'doc_preprocessor_res' in raw_result[0]:
211
  proc_res = raw_result[0]['doc_preprocessor_res']
212
  if 'output_img' in proc_res:
213
- print("Phát hiện ảnh đã qua xử lý hình học.")
214
  numpy_img = proc_res['output_img']
215
  target_image_for_drawing = Image.fromarray(numpy_img)
216
 
217
- # 4. Vẽ
218
  annotated_image = universal_draw(target_image_for_drawing, raw_result, FONT_PATH)
219
 
220
- # 5. Xử lý Text
221
  all_texts = deep_extract_text(raw_result)
222
  final_texts = clean_text_result(all_texts)
223
  text_output = "\n".join(final_texts) if final_texts else "Không tìm thấy văn bản."
224
 
225
- # Debug
226
  debug_str = str(raw_result)[:1000]
227
- lang_code_used = LANG_MAP.get(selected_lang, 'unknown')
228
- debug_info = f"Language: {selected_lang} (Code: {lang_code_used})\nData Preview:\n{debug_str}..."
229
 
230
  return annotated_image, text_output, debug_info
231
 
@@ -233,27 +214,30 @@ def predict(image, selected_lang):
233
  import traceback
234
  return image, f"Lỗi: {str(e)}", traceback.format_exc()
235
 
236
- # --- GIAO DIỆN (CẬP NHẬT: Thêm Dropdown) ---
237
- with gr.Blocks(title="PaddleOCR Multi-Lang Pro") as iface:
238
- gr.Markdown("## PaddleOCR Multi-Language - Precision Overlay")
239
 
240
  with gr.Row():
241
  with gr.Column():
242
  input_img = gr.Image(type="pil", label="Input Image")
243
 
244
- # Dropdown chọn ngôn ngữ
 
 
 
245
  lang_dropdown = gr.Dropdown(
246
- choices=list(LANG_MAP.keys()),
247
- value="Vietnamese (Tiếng Việt)", # Mặc định chọn Tiếng Việt
248
- label="Chọn Ngôn Ngữ (Select Language)",
249
- info="Chọn đúng ngôn ngữ để nhận diện dấu tốt nhất."
250
  )
251
 
252
  submit_btn = gr.Button("RUN OCR", variant="primary")
253
 
254
  with gr.Column():
255
  with gr.Tabs():
256
- with gr.TabItem("🖼️ Kết quả"):
257
  output_img = gr.Image(type="pil", label="Overlay Result")
258
  with gr.TabItem("📝 Văn bản"):
259
  output_txt = gr.Textbox(label="Text Content", lines=15)
@@ -262,7 +246,7 @@ with gr.Blocks(title="PaddleOCR Multi-Lang Pro") as iface:
262
 
263
  submit_btn.click(
264
  fn=predict,
265
- inputs=[input_img, lang_dropdown], # Truyền thêm input ngôn ngữ
266
  outputs=[output_img, output_txt, output_debug]
267
  )
268
 
 
17
  # Tắt log thừa
18
  logging.getLogger("ppocr").setLevel(logging.WARNING)
19
 
20
+ # --- QUẢN LÝ MODEL ĐA NGÔN NGỮ (THAY ĐỔI ĐỂ HỖ TRỢ THÊM NGÔN NGỮ) ---
21
  print("Đang khởi tạo hệ thống quản lý Model...")
22
 
23
+ # Cache để lưu các model đã tải, tránh load lại gây chậm
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
  OCR_ENGINES = {}
25
 
26
+ def get_ocr_model(lang_code):
27
+ # Nếu model ngôn ngữ này đã load rồi thì dùng luôn
 
 
 
28
  if lang_code in OCR_ENGINES:
29
  return OCR_ENGINES[lang_code]
30
 
31
+ print(f"Đang tải model ngôn ngữ: {lang_code}...")
32
  try:
33
  # Khởi tạo PaddleOCR với ngôn ngữ được chọn
34
+ # 'vi' hỗ trợ tốt Tiếng Việt và Latin. 'en' hỗ trợ chung Latin. 'ch' là Tiếng Trung.
35
  engine = PaddleOCR(
36
  use_textline_orientation=True,
37
  use_doc_orientation_classify=False,
 
41
  OCR_ENGINES[lang_code] = engine
42
  return engine
43
  except Exception as e:
44
+ print(f"Lỗi khởi tạo ngôn ngữ {lang_code}: {e}. Thử khởi tạo lại chế độ bản.")
45
+ # Fallback nếu tham số lỗi
46
  engine = PaddleOCR(lang='en')
47
  OCR_ENGINES[lang_code] = engine
48
  return engine
 
66
  def universal_draw(image, raw_data, font_path):
67
  if image is None: return image
68
 
69
+ # Đảm bảo image là PIL
70
  if isinstance(image, np.ndarray):
71
  image = Image.fromarray(image)
72
 
73
+ # Copy để vẽ
74
  canvas = image.copy()
75
  draw = ImageDraw.Draw(canvas)
76
 
 
80
  except:
81
  font = ImageFont.load_default()
82
 
83
+ # Hàm parse box
84
  def parse_box(b):
85
  try:
86
  if hasattr(b, 'tolist'): b = b.tolist()
 
92
 
93
  items_to_draw = []
94
 
95
+ # Logic tìm box/text
96
+ # Ưu tiên cấu trúc PaddleX: rec_texts + dt_polys
97
  processed = False
98
  if isinstance(raw_data, list) and len(raw_data) > 0 and isinstance(raw_data[0], dict):
99
  data_dict = raw_data[0]
 
107
  if box and txt: items_to_draw.append((box, txt))
108
  processed = True
109
 
110
+ # Fallback Logic
111
  if not processed:
112
  def hunt(data):
113
  if isinstance(data, dict):
 
127
  for item in data: hunt(item)
128
  hunt(raw_data)
129
 
130
+ # Vẽ
131
  for box, txt in items_to_draw:
132
  try:
133
+ # Vẽ khung đỏ
134
  draw.polygon(box, outline="red", width=3)
135
+ # Vẽ chữ
136
  txt_x, txt_y = box[0]
137
  if hasattr(draw, "textbbox"):
138
  text_bbox = draw.textbbox((txt_x, txt_y), txt, font=font, anchor="lb")
 
169
  cleaned.append(t)
170
  return cleaned
171
 
172
+ # --- MAIN PREDICT (CẬP NHẬT THAM SỐ LANG) ---
173
+ def predict(image, lang_code):
174
  if image is None: return None, "Chưa có ảnh.", "No Data"
175
 
176
  try:
177
+ # Lấy model tương ứng với ngôn ngữ đã chọn
178
+ current_ocr = get_ocr_model(lang_code)
179
 
180
+ # Chuẩn bị ảnh đầu vào
181
  original_pil = image.copy() if isinstance(image, Image.Image) else Image.fromarray(image).copy()
182
  image_np = np.array(image)
183
 
184
+ # 1. OCR (Sử dụng model đã chọn)
185
  raw_result = current_ocr.ocr(image_np)
186
 
187
+ # 2. XỬ LÝ ẢNH ĐỂ VẼ
188
  target_image_for_drawing = original_pil
189
+
190
+ # Kiểm tra xem Paddle có chỉnh sửa ảnh không
191
  if isinstance(raw_result, list) and len(raw_result) > 0 and isinstance(raw_result[0], dict):
192
  if 'doc_preprocessor_res' in raw_result[0]:
193
  proc_res = raw_result[0]['doc_preprocessor_res']
194
  if 'output_img' in proc_res:
195
+ print("Phát hiện ảnh đã qua xử lý hình học. Đang đồng bộ tọa độ...")
196
  numpy_img = proc_res['output_img']
197
  target_image_for_drawing = Image.fromarray(numpy_img)
198
 
199
+ # 3. Vẽ lên ảnh ĐÚNG (Target Image)
200
  annotated_image = universal_draw(target_image_for_drawing, raw_result, FONT_PATH)
201
 
202
+ # 4. Xử lý Text
203
  all_texts = deep_extract_text(raw_result)
204
  final_texts = clean_text_result(all_texts)
205
  text_output = "\n".join(final_texts) if final_texts else "Không tìm thấy văn bản."
206
 
207
+ # Debug Info
208
  debug_str = str(raw_result)[:1000]
209
+ debug_info = f"Language Mode: {lang_code}\nUsed Image Source: {'Preprocessed' if target_image_for_drawing != original_pil else 'Original'}\nData Preview:\n{debug_str}..."
 
210
 
211
  return annotated_image, text_output, debug_info
212
 
 
214
  import traceback
215
  return image, f"Lỗi: {str(e)}", traceback.format_exc()
216
 
217
+ # --- GIAO DIỆN (CẬP NHẬT DROPDOWN) ---
218
+ with gr.Blocks(title="PaddleOCR Multi-Language Overlay") as iface:
219
+ gr.Markdown("## PaddleOCR Multi-Language - High Precision Overlay")
220
 
221
  with gr.Row():
222
  with gr.Column():
223
  input_img = gr.Image(type="pil", label="Input Image")
224
 
225
+ # Thêm Dropdown chọn ngôn ngữ
226
+ # 'vi' bao phủ hầu hết các ngôn ngữ Latin + Tiếng Việt
227
+ # 'ch' cho tiếng Trung
228
+ # 'en', 'fr', 'de' cho các model chuyên biệt nếu cần
229
  lang_dropdown = gr.Dropdown(
230
+ choices=["vi", "ch", "en", "fr", "de", "es", "it", "pt", "ru", "ja", "ko"],
231
+ value="vi",
232
+ label="Chọn Ngôn Ngữ (Language)",
233
+ info="Chọn 'vi' cho Tiếng Việt & hầu hết ngôn ngữ Latin (Anh, Pháp, Đức...)."
234
  )
235
 
236
  submit_btn = gr.Button("RUN OCR", variant="primary")
237
 
238
  with gr.Column():
239
  with gr.Tabs():
240
+ with gr.TabItem("🖼️ Kết quả Khớp Tọa Độ"):
241
  output_img = gr.Image(type="pil", label="Overlay Result")
242
  with gr.TabItem("📝 Văn bản"):
243
  output_txt = gr.Textbox(label="Text Content", lines=15)
 
246
 
247
  submit_btn.click(
248
  fn=predict,
249
+ inputs=[input_img, lang_dropdown], # Truyền thêm lang_dropdown
250
  outputs=[output_img, output_txt, output_debug]
251
  )
252