langparse 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (101) hide show
  1. langparse/__init__.py +55 -0
  2. langparse/autoparser.py +25 -0
  3. langparse/chunkers/__init__.py +12 -0
  4. langparse/chunkers/blocks.py +151 -0
  5. langparse/chunkers/profiles.py +53 -0
  6. langparse/chunkers/registry.py +38 -0
  7. langparse/chunkers/semantic.py +242 -0
  8. langparse/chunkers/text.py +96 -0
  9. langparse/chunkers/workbook.py +942 -0
  10. langparse/cli.py +329 -0
  11. langparse/config.py +169 -0
  12. langparse/core/__init__.py +0 -0
  13. langparse/core/chunker.py +16 -0
  14. langparse/core/engine.py +37 -0
  15. langparse/core/parser.py +35 -0
  16. langparse/core/rendering.py +49 -0
  17. langparse/engines/__init__.py +1 -0
  18. langparse/engines/pdf/__init__.py +1 -0
  19. langparse/engines/pdf/deepdoc/__init__.py +55 -0
  20. langparse/engines/pdf/deepdoc/layout_recognizer.py +235 -0
  21. langparse/engines/pdf/deepdoc/model_loader.py +101 -0
  22. langparse/engines/pdf/deepdoc/ocr.py +641 -0
  23. langparse/engines/pdf/deepdoc/operators.py +684 -0
  24. langparse/engines/pdf/deepdoc/pdf_parser.py +1894 -0
  25. langparse/engines/pdf/deepdoc/postprocess.py +339 -0
  26. langparse/engines/pdf/deepdoc/recognizer.py +418 -0
  27. langparse/engines/pdf/deepdoc/rendering.py +210 -0
  28. langparse/engines/pdf/deepdoc/table_structure_recognizer.py +559 -0
  29. langparse/engines/pdf/deepdoc/tokenizer.py +30 -0
  30. langparse/engines/pdf/deepdoc/utils.py +36 -0
  31. langparse/engines/pdf/deepdoc_engine.py +164 -0
  32. langparse/engines/pdf/mineru.py +259 -0
  33. langparse/engines/pdf/mineru_client.py +318 -0
  34. langparse/engines/pdf/mineru_service.py +225 -0
  35. langparse/engines/pdf/ocr.py +101 -0
  36. langparse/engines/pdf/other.py +20 -0
  37. langparse/engines/pdf/simple.py +134 -0
  38. langparse/engines/pdf/vision_llm.py +27 -0
  39. langparse/errors.py +70 -0
  40. langparse/logging.py +27 -0
  41. langparse/metrics.py +129 -0
  42. langparse/parsers/__init__.py +0 -0
  43. langparse/parsers/docx_parser.py +114 -0
  44. langparse/parsers/excel_parser.py +220 -0
  45. langparse/parsers/markdown_parser.py +34 -0
  46. langparse/parsers/pdf_parser.py +31 -0
  47. langparse/parsers/registry.py +48 -0
  48. langparse/parsers/sniff.py +72 -0
  49. langparse/progress.py +77 -0
  50. langparse/py.typed +0 -0
  51. langparse/services/__init__.py +11 -0
  52. langparse/services/batch_service.py +339 -0
  53. langparse/services/benchmark_service.py +202 -0
  54. langparse/services/fidelity.py +154 -0
  55. langparse/services/output_paths.py +86 -0
  56. langparse/services/parse_service.py +523 -0
  57. langparse/services/quality.py +65 -0
  58. langparse/services/workbook_ambiguity_benchmark.py +563 -0
  59. langparse/services/workbook_quality_benchmark.py +230 -0
  60. langparse/types.py +97 -0
  61. langparse/workbooks/__init__.py +103 -0
  62. langparse/workbooks/adapters.py +474 -0
  63. langparse/workbooks/assembly.py +993 -0
  64. langparse/workbooks/blocks.py +209 -0
  65. langparse/workbooks/bundle-v1.schema.json +71 -0
  66. langparse/workbooks/bundle.py +341 -0
  67. langparse/workbooks/classification.py +393 -0
  68. langparse/workbooks/continuation.py +577 -0
  69. langparse/workbooks/evaluation/__init__.py +45 -0
  70. langparse/workbooks/evaluation/evaluator.py +381 -0
  71. langparse/workbooks/evaluation/schema.py +419 -0
  72. langparse/workbooks/labels.py +14 -0
  73. langparse/workbooks/lineage.py +117 -0
  74. langparse/workbooks/modeling/__init__.py +52 -0
  75. langparse/workbooks/modeling/cache.py +20 -0
  76. langparse/workbooks/modeling/config.py +87 -0
  77. langparse/workbooks/modeling/contract.py +628 -0
  78. langparse/workbooks/modeling/disambiguation.py +800 -0
  79. langparse/workbooks/modeling/openai_adapter.py +192 -0
  80. langparse/workbooks/modeling/policy.py +79 -0
  81. langparse/workbooks/modeling/ports.py +44 -0
  82. langparse/workbooks/modeling/pricing.py +17 -0
  83. langparse/workbooks/modeling/types.py +251 -0
  84. langparse/workbooks/objects.py +229 -0
  85. langparse/workbooks/quality/__init__.py +23 -0
  86. langparse/workbooks/quality/bundle.py +53 -0
  87. langparse/workbooks/quality/evaluator.py +266 -0
  88. langparse/workbooks/quality/facts.py +142 -0
  89. langparse/workbooks/quality/schema.py +462 -0
  90. langparse/workbooks/reference_types.py +73 -0
  91. langparse/workbooks/references.py +178 -0
  92. langparse/workbooks/regions.py +932 -0
  93. langparse/workbooks/rendering.py +222 -0
  94. langparse/workbooks/tables.py +477 -0
  95. langparse/workbooks/types.py +257 -0
  96. langparse-0.1.0.dist-info/METADATA +790 -0
  97. langparse-0.1.0.dist-info/RECORD +101 -0
  98. langparse-0.1.0.dist-info/WHEEL +5 -0
  99. langparse-0.1.0.dist-info/entry_points.txt +2 -0
  100. langparse-0.1.0.dist-info/licenses/LICENSE +192 -0
  101. langparse-0.1.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,641 @@
1
+ #
2
+ # Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ #
16
+ import gc
17
+ import logging
18
+ import copy
19
+ import time
20
+ import os
21
+
22
+ from .model_loader import default_model_dir
23
+ from .operators import * # noqa: F403
24
+ from . import operators
25
+ import math
26
+ import numpy as np
27
+ import cv2
28
+ import onnxruntime as ort
29
+
30
+ from .postprocess import build_post_process
31
+
32
+ loaded_models = {}
33
+
34
+
35
+ def transform(data, ops=None):
36
+ """transform"""
37
+ if ops is None:
38
+ ops = []
39
+ for op in ops:
40
+ data = op(data)
41
+ if data is None:
42
+ return None
43
+ return data
44
+
45
+
46
+ def create_operators(op_param_list, global_config=None):
47
+ """
48
+ create operators based on the config
49
+
50
+ Args:
51
+ params(list): a dict list, used to create some operators
52
+ """
53
+ assert isinstance(op_param_list, list), "operator config should be a list"
54
+ ops = []
55
+ for operator in op_param_list:
56
+ assert isinstance(operator, dict) and len(operator) == 1, "yaml format error"
57
+ op_name = list(operator)[0]
58
+ param = {} if operator[op_name] is None else operator[op_name]
59
+ if global_config is not None:
60
+ param.update(global_config)
61
+ op = getattr(operators, op_name)(**param)
62
+ ops.append(op)
63
+ return ops
64
+
65
+
66
+ def load_model(model_dir, nm, device_id: int | None = None):
67
+ model_file_path = os.path.join(model_dir, nm + ".onnx")
68
+ model_cached_tag = model_file_path + str(device_id) if device_id is not None else model_file_path
69
+
70
+ global loaded_models
71
+ loaded_model = loaded_models.get(model_cached_tag)
72
+ if loaded_model:
73
+ logging.info(f"load_model {model_file_path} reuses cached model")
74
+ return loaded_model
75
+
76
+ if not os.path.exists(model_file_path):
77
+ raise ValueError("not find model file path {}".format(model_file_path))
78
+
79
+ def cuda_is_available():
80
+ try:
81
+ import torch
82
+ except ImportError:
83
+ return False
84
+
85
+ target_id = 0 if device_id is None else device_id
86
+ return bool(torch.cuda.is_available() and torch.cuda.device_count() > target_id)
87
+
88
+ options = ort.SessionOptions()
89
+ options.enable_cpu_mem_arena = False
90
+ options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
91
+ # Prevent CPU oversubscription by allowing explicit thread control in multi-worker environments
92
+ options.intra_op_num_threads = int(os.environ.get("OCR_INTRA_OP_NUM_THREADS", "2"))
93
+ options.inter_op_num_threads = int(os.environ.get("OCR_INTER_OP_NUM_THREADS", "2"))
94
+
95
+ # https://github.com/microsoft/onnxruntime/issues/9509#issuecomment-951546580
96
+ # Shrink GPU memory after execution
97
+ run_options = ort.RunOptions()
98
+ if cuda_is_available():
99
+ gpu_mem_limit_mb = int(os.environ.get("OCR_GPU_MEM_LIMIT_MB", "2048"))
100
+ arena_strategy = os.environ.get("OCR_ARENA_EXTEND_STRATEGY", "kNextPowerOfTwo")
101
+ provider_device_id = 0 if device_id is None else device_id
102
+ cuda_provider_options = {
103
+ "device_id": provider_device_id, # Use specific GPU
104
+ "gpu_mem_limit": max(gpu_mem_limit_mb, 0) * 1024 * 1024,
105
+ "arena_extend_strategy": arena_strategy, # gpu memory allocation strategy
106
+ }
107
+ sess = ort.InferenceSession(model_file_path, options=options, providers=["CUDAExecutionProvider"], provider_options=[cuda_provider_options])
108
+ # Explicit arena shrinkage for GPU to release VRAM back to the system after each run
109
+ if os.environ.get("OCR_GPUMEM_ARENA_SHRINKAGE") == "1":
110
+ run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", f"gpu:{provider_device_id}")
111
+ logging.info(f"load_model {model_file_path} enabled GPU memory arena shrinkage on device {provider_device_id}")
112
+ logging.info(f"load_model {model_file_path} uses GPU (device {provider_device_id}, gpu_mem_limit={cuda_provider_options['gpu_mem_limit']}, arena_strategy={arena_strategy})")
113
+ else:
114
+ sess = ort.InferenceSession(model_file_path, options=options, providers=["CPUExecutionProvider"])
115
+ run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "cpu")
116
+ logging.info(f"load_model {model_file_path} uses CPU")
117
+ loaded_model = (sess, run_options)
118
+ loaded_models[model_cached_tag] = loaded_model
119
+ return loaded_model
120
+
121
+
122
+ class TextRecognizer:
123
+ def __init__(self, model_dir, device_id: int | None = None):
124
+ self.rec_image_shape = [int(v) for v in "3, 48, 320".split(",")]
125
+ self.rec_batch_num = 16
126
+ postprocess_params = {"name": "CTCLabelDecode", "character_dict_path": os.path.join(model_dir, "ocr.res"), "use_space_char": True}
127
+ self.postprocess_op = build_post_process(postprocess_params)
128
+ self.predictor, self.run_options = load_model(model_dir, "rec", device_id)
129
+ self.input_tensor = self.predictor.get_inputs()[0]
130
+
131
+ def resize_norm_img(self, img, max_wh_ratio):
132
+ imgC, imgH, imgW = self.rec_image_shape
133
+
134
+ assert imgC == img.shape[2]
135
+ imgW = int((imgH * max_wh_ratio))
136
+ w = self.input_tensor.shape[3:][0]
137
+ if isinstance(w, str):
138
+ pass
139
+ elif w is not None and w > 0:
140
+ imgW = w
141
+ h, w = img.shape[:2]
142
+ ratio = w / float(h)
143
+ if math.ceil(imgH * ratio) > imgW:
144
+ resized_w = imgW
145
+ else:
146
+ resized_w = int(math.ceil(imgH * ratio))
147
+
148
+ resized_image = cv2.resize(img, (resized_w, imgH))
149
+ resized_image = resized_image.astype("float32")
150
+ resized_image = resized_image.transpose((2, 0, 1)) / 255
151
+ resized_image -= 0.5
152
+ resized_image /= 0.5
153
+ padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
154
+ padding_im[:, :, 0:resized_w] = resized_image
155
+ return padding_im
156
+
157
+ def resize_norm_img_vl(self, img, image_shape):
158
+
159
+ imgC, imgH, imgW = image_shape
160
+ img = img[:, :, ::-1] # bgr2rgb
161
+ resized_image = cv2.resize(img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
162
+ resized_image = resized_image.astype("float32")
163
+ resized_image = resized_image.transpose((2, 0, 1)) / 255
164
+ return resized_image
165
+
166
+ def resize_norm_img_srn(self, img, image_shape):
167
+ imgC, imgH, imgW = image_shape
168
+
169
+ img_black = np.zeros((imgH, imgW))
170
+ im_hei = img.shape[0]
171
+ im_wid = img.shape[1]
172
+
173
+ if im_wid <= im_hei * 1:
174
+ img_new = cv2.resize(img, (imgH * 1, imgH))
175
+ elif im_wid <= im_hei * 2:
176
+ img_new = cv2.resize(img, (imgH * 2, imgH))
177
+ elif im_wid <= im_hei * 3:
178
+ img_new = cv2.resize(img, (imgH * 3, imgH))
179
+ else:
180
+ img_new = cv2.resize(img, (imgW, imgH))
181
+
182
+ img_np = np.asarray(img_new)
183
+ img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
184
+ img_black[:, 0 : img_np.shape[1]] = img_np
185
+ img_black = img_black[:, :, np.newaxis]
186
+
187
+ row, col, c = img_black.shape
188
+ c = 1
189
+
190
+ return np.reshape(img_black, (c, row, col)).astype(np.float32)
191
+
192
+ def srn_other_inputs(self, image_shape, num_heads, max_text_length):
193
+
194
+ imgC, imgH, imgW = image_shape
195
+ feature_dim = int((imgH / 8) * (imgW / 8))
196
+
197
+ encoder_word_pos = np.array(range(0, feature_dim)).reshape((feature_dim, 1)).astype("int64")
198
+ gsrm_word_pos = np.array(range(0, max_text_length)).reshape((max_text_length, 1)).astype("int64")
199
+
200
+ gsrm_attn_bias_data = np.ones((1, max_text_length, max_text_length))
201
+ gsrm_slf_attn_bias1 = np.triu(gsrm_attn_bias_data, 1).reshape([-1, 1, max_text_length, max_text_length])
202
+ gsrm_slf_attn_bias1 = np.tile(gsrm_slf_attn_bias1, [1, num_heads, 1, 1]).astype("float32") * [-1e9]
203
+
204
+ gsrm_slf_attn_bias2 = np.tril(gsrm_attn_bias_data, -1).reshape([-1, 1, max_text_length, max_text_length])
205
+ gsrm_slf_attn_bias2 = np.tile(gsrm_slf_attn_bias2, [1, num_heads, 1, 1]).astype("float32") * [-1e9]
206
+
207
+ encoder_word_pos = encoder_word_pos[np.newaxis, :]
208
+ gsrm_word_pos = gsrm_word_pos[np.newaxis, :]
209
+
210
+ return [encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2]
211
+
212
+ def process_image_srn(self, img, image_shape, num_heads, max_text_length):
213
+ norm_img = self.resize_norm_img_srn(img, image_shape)
214
+ norm_img = norm_img[np.newaxis, :]
215
+
216
+ [encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2] = self.srn_other_inputs(image_shape, num_heads, max_text_length)
217
+
218
+ gsrm_slf_attn_bias1 = gsrm_slf_attn_bias1.astype(np.float32)
219
+ gsrm_slf_attn_bias2 = gsrm_slf_attn_bias2.astype(np.float32)
220
+ encoder_word_pos = encoder_word_pos.astype(np.int64)
221
+ gsrm_word_pos = gsrm_word_pos.astype(np.int64)
222
+
223
+ return (norm_img, encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2)
224
+
225
+ def resize_norm_img_sar(self, img, image_shape, width_downsample_ratio=0.25):
226
+ imgC, imgH, imgW_min, imgW_max = image_shape
227
+ h = img.shape[0]
228
+ w = img.shape[1]
229
+ valid_ratio = 1.0
230
+ # make sure new_width is an integral multiple of width_divisor.
231
+ width_divisor = int(1 / width_downsample_ratio)
232
+ # resize
233
+ ratio = w / float(h)
234
+ resize_w = math.ceil(imgH * ratio)
235
+ if resize_w % width_divisor != 0:
236
+ resize_w = round(resize_w / width_divisor) * width_divisor
237
+ if imgW_min is not None:
238
+ resize_w = max(imgW_min, resize_w)
239
+ if imgW_max is not None:
240
+ valid_ratio = min(1.0, 1.0 * resize_w / imgW_max)
241
+ resize_w = min(imgW_max, resize_w)
242
+ resized_image = cv2.resize(img, (resize_w, imgH))
243
+ resized_image = resized_image.astype("float32")
244
+ # norm
245
+ if image_shape[0] == 1:
246
+ resized_image = resized_image / 255
247
+ resized_image = resized_image[np.newaxis, :]
248
+ else:
249
+ resized_image = resized_image.transpose((2, 0, 1)) / 255
250
+ resized_image -= 0.5
251
+ resized_image /= 0.5
252
+ resize_shape = resized_image.shape
253
+ padding_im = -1.0 * np.ones((imgC, imgH, imgW_max), dtype=np.float32)
254
+ padding_im[:, :, 0:resize_w] = resized_image
255
+ pad_shape = padding_im.shape
256
+
257
+ return padding_im, resize_shape, pad_shape, valid_ratio
258
+
259
+ def resize_norm_img_spin(self, img):
260
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
261
+ # return padding_im
262
+ img = cv2.resize(img, tuple([100, 32]), cv2.INTER_CUBIC)
263
+ img = np.array(img, np.float32)
264
+ img = np.expand_dims(img, -1)
265
+ img = img.transpose((2, 0, 1))
266
+ mean = [127.5]
267
+ std = [127.5]
268
+ mean = np.array(mean, dtype=np.float32)
269
+ std = np.array(std, dtype=np.float32)
270
+ mean = np.float32(mean.reshape(1, -1))
271
+ stdinv = 1 / np.float32(std.reshape(1, -1))
272
+ img -= mean
273
+ img *= stdinv
274
+ return img
275
+
276
+ def resize_norm_img_svtr(self, img, image_shape):
277
+
278
+ imgC, imgH, imgW = image_shape
279
+ resized_image = cv2.resize(img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
280
+ resized_image = resized_image.astype("float32")
281
+ resized_image = resized_image.transpose((2, 0, 1)) / 255
282
+ resized_image -= 0.5
283
+ resized_image /= 0.5
284
+ return resized_image
285
+
286
+ def resize_norm_img_abinet(self, img, image_shape):
287
+
288
+ imgC, imgH, imgW = image_shape
289
+
290
+ resized_image = cv2.resize(img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
291
+ resized_image = resized_image.astype("float32")
292
+ resized_image = resized_image / 255.0
293
+
294
+ mean = np.array([0.485, 0.456, 0.406])
295
+ std = np.array([0.229, 0.224, 0.225])
296
+ resized_image = (resized_image - mean[None, None, ...]) / std[None, None, ...]
297
+ resized_image = resized_image.transpose((2, 0, 1))
298
+ resized_image = resized_image.astype("float32")
299
+
300
+ return resized_image
301
+
302
+ def norm_img_can(self, img, image_shape):
303
+
304
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # CAN only predict gray scale image
305
+
306
+ if self.rec_image_shape[0] == 1:
307
+ h, w = img.shape
308
+ _, imgH, imgW = self.rec_image_shape
309
+ if h < imgH or w < imgW:
310
+ padding_h = max(imgH - h, 0)
311
+ padding_w = max(imgW - w, 0)
312
+ img_padded = np.pad(img, ((0, padding_h), (0, padding_w)), "constant", constant_values=(255))
313
+ img = img_padded
314
+
315
+ img = np.expand_dims(img, 0) / 255.0 # h,w,c -> c,h,w
316
+ img = img.astype("float32")
317
+
318
+ return img
319
+
320
+ def close(self):
321
+ # close session and release manually
322
+ logging.info("Close text recognizer.")
323
+ if hasattr(self, "predictor"):
324
+ del self.predictor
325
+ gc.collect()
326
+
327
+ def __call__(self, img_list):
328
+ img_num = len(img_list)
329
+ # Calculate the aspect ratio of all text bars
330
+ width_list = []
331
+ for img in img_list:
332
+ width_list.append(img.shape[1] / float(img.shape[0]))
333
+ # Sorting can speed up the recognition process
334
+ indices = np.argsort(np.array(width_list))
335
+ rec_res = [["", 0.0]] * img_num
336
+ batch_num = self.rec_batch_num
337
+ st = time.time()
338
+
339
+ for beg_img_no in range(0, img_num, batch_num):
340
+ end_img_no = min(img_num, beg_img_no + batch_num)
341
+ norm_img_batch = []
342
+ imgC, imgH, imgW = self.rec_image_shape[:3]
343
+ max_wh_ratio = imgW / imgH
344
+ # max_wh_ratio = 0
345
+ for ino in range(beg_img_no, end_img_no):
346
+ h, w = img_list[indices[ino]].shape[0:2]
347
+ wh_ratio = w * 1.0 / h
348
+ max_wh_ratio = max(max_wh_ratio, wh_ratio)
349
+ for ino in range(beg_img_no, end_img_no):
350
+ norm_img = self.resize_norm_img(img_list[indices[ino]], max_wh_ratio)
351
+ norm_img = norm_img[np.newaxis, :]
352
+ norm_img_batch.append(norm_img)
353
+ norm_img_batch = np.concatenate(norm_img_batch)
354
+ norm_img_batch = norm_img_batch.copy()
355
+
356
+ input_dict = {}
357
+ input_dict[self.input_tensor.name] = norm_img_batch
358
+ for i in range(100000):
359
+ try:
360
+ outputs = self.predictor.run(None, input_dict, self.run_options)
361
+ break
362
+ except Exception as e:
363
+ if i >= 3:
364
+ raise e
365
+ time.sleep(5)
366
+ preds = outputs[0]
367
+ rec_result = self.postprocess_op(preds)
368
+ for rno in range(len(rec_result)):
369
+ rec_res[indices[beg_img_no + rno]] = rec_result[rno]
370
+
371
+ return rec_res, time.time() - st
372
+
373
+ def __del__(self):
374
+ self.close()
375
+
376
+
377
+ class TextDetector:
378
+ def __init__(self, model_dir, device_id: int | None = None):
379
+ pre_process_list = [
380
+ {
381
+ "DetResizeForTest": {
382
+ "limit_side_len": 960,
383
+ "limit_type": "max",
384
+ }
385
+ },
386
+ {"NormalizeImage": {"std": [0.229, 0.224, 0.225], "mean": [0.485, 0.456, 0.406], "scale": "1./255.", "order": "hwc"}},
387
+ {"ToCHWImage": None},
388
+ {"KeepKeys": {"keep_keys": ["image", "shape"]}},
389
+ ]
390
+ postprocess_params = {"name": "DBPostProcess", "thresh": 0.3, "box_thresh": 0.5, "max_candidates": 1000, "unclip_ratio": 1.5, "use_dilation": False, "score_mode": "fast", "box_type": "quad"}
391
+
392
+ self.postprocess_op = build_post_process(postprocess_params)
393
+ self.predictor, self.run_options = load_model(model_dir, "det", device_id)
394
+ self.input_tensor = self.predictor.get_inputs()[0]
395
+
396
+ img_h, img_w = self.input_tensor.shape[2:]
397
+ if isinstance(img_h, str) or isinstance(img_w, str):
398
+ pass
399
+ elif img_h is not None and img_w is not None and img_h > 0 and img_w > 0:
400
+ pre_process_list[0] = {"DetResizeForTest": {"image_shape": [img_h, img_w]}}
401
+ self.preprocess_op = create_operators(pre_process_list)
402
+
403
+ def order_points_clockwise(self, pts):
404
+ rect = np.zeros((4, 2), dtype="float32")
405
+ s = pts.sum(axis=1)
406
+ rect[0] = pts[np.argmin(s)]
407
+ rect[2] = pts[np.argmax(s)]
408
+ tmp = np.delete(pts, (np.argmin(s), np.argmax(s)), axis=0)
409
+ diff = np.diff(np.array(tmp), axis=1)
410
+ rect[1] = tmp[np.argmin(diff)]
411
+ rect[3] = tmp[np.argmax(diff)]
412
+ return rect
413
+
414
+ def clip_det_res(self, points, img_height, img_width):
415
+ for pno in range(points.shape[0]):
416
+ points[pno, 0] = int(min(max(points[pno, 0], 0), img_width - 1))
417
+ points[pno, 1] = int(min(max(points[pno, 1], 0), img_height - 1))
418
+ return points
419
+
420
+ def filter_tag_det_res(self, dt_boxes, image_shape):
421
+ img_height, img_width = image_shape[0:2]
422
+ dt_boxes_new = []
423
+ for box in dt_boxes:
424
+ if isinstance(box, list):
425
+ box = np.array(box)
426
+ box = self.order_points_clockwise(box)
427
+ box = self.clip_det_res(box, img_height, img_width)
428
+ rect_width = int(np.linalg.norm(box[0] - box[1]))
429
+ rect_height = int(np.linalg.norm(box[0] - box[3]))
430
+ if rect_width <= 3 or rect_height <= 3:
431
+ continue
432
+ dt_boxes_new.append(box)
433
+ dt_boxes = np.array(dt_boxes_new)
434
+ return dt_boxes
435
+
436
+ def filter_tag_det_res_only_clip(self, dt_boxes, image_shape):
437
+ img_height, img_width = image_shape[0:2]
438
+ dt_boxes_new = []
439
+ for box in dt_boxes:
440
+ if isinstance(box, list):
441
+ box = np.array(box)
442
+ box = self.clip_det_res(box, img_height, img_width)
443
+ dt_boxes_new.append(box)
444
+ dt_boxes = np.array(dt_boxes_new)
445
+ return dt_boxes
446
+
447
+ def close(self):
448
+ logging.info("Close text detector.")
449
+ if hasattr(self, "predictor"):
450
+ del self.predictor
451
+ gc.collect()
452
+
453
+ def __call__(self, img):
454
+ ori_im = img.copy()
455
+ data = {"image": img}
456
+
457
+ st = time.time()
458
+ data = transform(data, self.preprocess_op)
459
+ img, shape_list = data
460
+ if img is None:
461
+ return None, 0
462
+ img = np.expand_dims(img, axis=0)
463
+ shape_list = np.expand_dims(shape_list, axis=0)
464
+ img = img.copy()
465
+ input_dict = {}
466
+ input_dict[self.input_tensor.name] = img
467
+ for i in range(100000):
468
+ try:
469
+ outputs = self.predictor.run(None, input_dict, self.run_options)
470
+ break
471
+ except Exception as e:
472
+ if i >= 3:
473
+ raise e
474
+ time.sleep(5)
475
+
476
+ post_result = self.postprocess_op({"maps": outputs[0]}, shape_list)
477
+ dt_boxes = post_result[0]["points"]
478
+ dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape)
479
+
480
+ return dt_boxes, time.time() - st
481
+
482
+ def __del__(self):
483
+ self.close()
484
+
485
+
486
+ class OCR:
487
+ def __init__(self, model_dir=None):
488
+ model_dir = model_dir or str(default_model_dir())
489
+ self.text_detector = [TextDetector(model_dir)]
490
+ self.text_recognizer = [TextRecognizer(model_dir)]
491
+ self.drop_score = 0.5
492
+ self.crop_image_res_index = 0
493
+
494
+ def get_rotate_crop_image(self, img, points):
495
+ """
496
+ img_height, img_width = img.shape[0:2]
497
+ left = int(np.min(points[:, 0]))
498
+ right = int(np.max(points[:, 0]))
499
+ top = int(np.min(points[:, 1]))
500
+ bottom = int(np.max(points[:, 1]))
501
+ img_crop = img[top:bottom, left:right, :].copy()
502
+ points[:, 0] = points[:, 0] - left
503
+ points[:, 1] = points[:, 1] - top
504
+ """
505
+ assert len(points) == 4, "shape of points must be 4*2"
506
+ img_crop_width = int(max(np.linalg.norm(points[0] - points[1]), np.linalg.norm(points[2] - points[3])))
507
+ img_crop_height = int(max(np.linalg.norm(points[0] - points[3]), np.linalg.norm(points[1] - points[2])))
508
+ pts_std = np.float32([[0, 0], [img_crop_width, 0], [img_crop_width, img_crop_height], [0, img_crop_height]])
509
+ M = cv2.getPerspectiveTransform(points, pts_std)
510
+ dst_img = cv2.warpPerspective(img, M, (img_crop_width, img_crop_height), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC)
511
+ dst_img_height, dst_img_width = dst_img.shape[0:2]
512
+ if dst_img_height * 1.0 / dst_img_width >= 1.5:
513
+ # Try original orientation
514
+ rec_result = self.text_recognizer[0]([dst_img])
515
+ text, score = rec_result[0][0]
516
+ best_score = score
517
+ best_img = dst_img
518
+
519
+ # Try clockwise 90° rotation
520
+ rotated_cw = np.rot90(dst_img, k=3)
521
+ rec_result = self.text_recognizer[0]([rotated_cw])
522
+ rotated_cw_text, rotated_cw_score = rec_result[0][0]
523
+ if rotated_cw_score > best_score:
524
+ best_score = rotated_cw_score
525
+ best_img = rotated_cw
526
+
527
+ # Try counter-clockwise 90° rotation
528
+ rotated_ccw = np.rot90(dst_img, k=1)
529
+ rec_result = self.text_recognizer[0]([rotated_ccw])
530
+ rotated_ccw_text, rotated_ccw_score = rec_result[0][0]
531
+ if rotated_ccw_score > best_score:
532
+ best_img = rotated_ccw
533
+
534
+ # Use the best image
535
+ dst_img = best_img
536
+ return dst_img
537
+
538
+ def sorted_boxes(self, dt_boxes):
539
+ """
540
+ Sort text boxes in order from top to bottom, left to right
541
+ args:
542
+ dt_boxes(array):detected text boxes with shape [4, 2]
543
+ return:
544
+ sorted boxes(array) with shape [4, 2]
545
+ """
546
+ num_boxes = dt_boxes.shape[0]
547
+ sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
548
+ _boxes = list(sorted_boxes)
549
+
550
+ for i in range(num_boxes - 1):
551
+ for j in range(i, -1, -1):
552
+ if abs(_boxes[j + 1][0][1] - _boxes[j][0][1]) < 10 and (_boxes[j + 1][0][0] < _boxes[j][0][0]):
553
+ tmp = _boxes[j]
554
+ _boxes[j] = _boxes[j + 1]
555
+ _boxes[j + 1] = tmp
556
+ else:
557
+ break
558
+ return _boxes
559
+
560
+ def detect(self, img, device_id: int | None = None):
561
+ if device_id is None:
562
+ device_id = 0
563
+
564
+ if img is None:
565
+ return None
566
+
567
+ dt_boxes, _ = self.text_detector[device_id](img)
568
+
569
+ if dt_boxes is None:
570
+ return None
571
+
572
+ return zip(self.sorted_boxes(dt_boxes), [("", 0) for _ in range(len(dt_boxes))])
573
+
574
+ def recognize(self, ori_im, box, device_id: int | None = None):
575
+ if device_id is None:
576
+ device_id = 0
577
+
578
+ img_crop = self.get_rotate_crop_image(ori_im, box)
579
+
580
+ rec_res, elapse = self.text_recognizer[device_id]([img_crop])
581
+ text, score = rec_res[0]
582
+ if score < self.drop_score:
583
+ return ""
584
+ return text
585
+
586
+ def recognize_batch(self, img_list, device_id: int | None = None):
587
+ if device_id is None:
588
+ device_id = 0
589
+ rec_res, elapse = self.text_recognizer[device_id](img_list)
590
+ texts = []
591
+ for i in range(len(rec_res)):
592
+ text, score = rec_res[i]
593
+ if score < self.drop_score:
594
+ text = ""
595
+ texts.append(text)
596
+ return texts
597
+
598
+ def __call__(self, img, device_id=0, cls=True):
599
+ time_dict = {"det": 0, "rec": 0, "cls": 0, "all": 0}
600
+ if device_id is None:
601
+ device_id = 0
602
+
603
+ if img is None:
604
+ return None, None, time_dict
605
+
606
+ start = time.time()
607
+ ori_im = img.copy()
608
+ dt_boxes, elapse = self.text_detector[device_id](img)
609
+ time_dict["det"] = elapse
610
+
611
+ if dt_boxes is None:
612
+ end = time.time()
613
+ time_dict["all"] = end - start
614
+ return None, None, time_dict
615
+
616
+ img_crop_list = []
617
+
618
+ dt_boxes = self.sorted_boxes(dt_boxes)
619
+
620
+ for bno in range(len(dt_boxes)):
621
+ tmp_box = copy.deepcopy(dt_boxes[bno])
622
+ img_crop = self.get_rotate_crop_image(ori_im, tmp_box)
623
+ img_crop_list.append(img_crop)
624
+
625
+ rec_res, elapse = self.text_recognizer[device_id](img_crop_list)
626
+
627
+ time_dict["rec"] = elapse
628
+
629
+ filter_boxes, filter_rec_res = [], []
630
+ for box, rec_result in zip(dt_boxes, rec_res):
631
+ text, score = rec_result
632
+ if score >= self.drop_score:
633
+ filter_boxes.append(box)
634
+ filter_rec_res.append(rec_result)
635
+ end = time.time()
636
+ time_dict["all"] = end - start
637
+
638
+ # for bno in range(len(img_crop_list)):
639
+ # print(f"{bno}, {rec_res[bno]}")
640
+
641
+ return list(zip([a.tolist() for a in filter_boxes], filter_rec_res))