philcuriosity1024 commited on
Commit
6d39ea5
·
verified ·
1 Parent(s): 6cd8bf9

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -33,3 +33,12 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
 
 
 
 
 
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ axmodel/ax650/cls_npu1.axmodel filter=lfs diff=lfs merge=lfs -text
37
+ axmodel/ax650/cls_npu3.axmodel filter=lfs diff=lfs merge=lfs -text
38
+ axmodel/ax650/det_npu1.axmodel filter=lfs diff=lfs merge=lfs -text
39
+ axmodel/ax650/det_npu3.axmodel filter=lfs diff=lfs merge=lfs -text
40
+ axmodel/ax650/rec_npu1.axmodel filter=lfs diff=lfs merge=lfs -text
41
+ axmodel/ax650/rec_npu3.axmodel filter=lfs diff=lfs merge=lfs -text
42
+ fonts/simfang.ttf filter=lfs diff=lfs merge=lfs -text
43
+ res-ax.jpg filter=lfs diff=lfs merge=lfs -text
44
+ res-onnx.jpg filter=lfs diff=lfs merge=lfs -text
11.jpg ADDED
README-zh.md ADDED
@@ -0,0 +1,186 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ language:
4
+ - zh
5
+ base_model:
6
+ - PP-OCRv6_mobile_det
7
+ - PP-OCRv6_mobile_rec
8
+ - PP-LCNet_x0_25_textline_ori
9
+ pipeline_tag: text-classification
10
+ tags:
11
+ - OCR
12
+ - paddle
13
+ - PPOCRv6
14
+ - axera
15
+ ---
16
+
17
+ # PPOCR_v6
18
+ > [English](./README.md) | 中文
19
+
20
+ 本版本 PPOCR_v6 已转换为可在爱芯 NPU 上运行的模型,采用 **w8a16** 量化。
21
+
22
+ ## 转换工具链接
23
+
24
+ 如果您对模型转换感兴趣,可以通过以下链接导出 axmodel:
25
+
26
+ - [ax-samples-github](https://github.com/AXERA-TECH/ax-samples),其他好玩的样例
27
+
28
+ - [Pulsar2 文档,ONNX 转 axmodel 方法](https://pulsar2-docs.readthedocs.io/en/latest/pulsar2/introduction.html)
29
+
30
+ ## 支持平台
31
+
32
+ - AX650
33
+ - [M4N-Dock(爱芯派Pro)](https://wiki.sipeed.com/hardware/zh/maixIV/m4ndock/m4ndock.html)
34
+ - [M.2 加速卡](https://axcl-docs.readthedocs.io/zh-cn/latest/doc_guide_hardware.html)
35
+ - AX630C
36
+ - [爱芯派2](https://axera-pi-2-docs-cn.readthedocs.io/zh-cn/latest/index.html)
37
+ - [Module-LLM](https://docs.m5stack.com/zh_CN/module/Module-LLM)
38
+ - [LLM630 Compute Kit](https://docs.m5stack.com/zh_CN/core/LLM630%20Compute%20Kit)
39
+ - AX615
40
+ - [AX615 IPC SoC](https://www.axera-tech.com/zh-hans/product/2956.html)
41
+
42
+
43
+ ### 性能测试
44
+
45
+ |芯片|模型|npu_mode|耗时(ms)|
46
+ |--|--|--|--|
47
+ ||PP-OCRv6_small_det|NPU1|42.158|
48
+ ||PP-OCRv6_small_det|NPU3|23.837|
49
+ |AX650|PP-LCNet_x0_25_textline_ori|NPU1|0.294|
50
+ ||PP-LCNet_x0_25_textline_ori|NPU3|0.172|
51
+ ||PP-OCRv6_small_rec|NPU1|2.473|
52
+ ||PP-OCRv6_small_rec|NPU3|1.073|
53
+
54
+ 测试方案:
55
+
56
+ ``` bash
57
+ ax_run_model -w 10 -r 100 -m xx.axmodel
58
+ ```
59
+
60
+ 识别和检测的onnx模型来源:[PaddlePaddle/PP-OCRv6](https://huggingface.co/collections/PaddlePaddle/pp-ocrv6) 中 [small-det-onnx](https://huggingface.co/PaddlePaddle/PP-OCRv6_small_det_onnx) 和 [small-rec-onnx](https://huggingface.co/PaddlePaddle/PP-OCRv6_small_rec_onnx)
61
+ 方向分类器:[AXERA-TECH/PPOCR_v5](https://huggingface.co/AXERA-TECH/PPOCR_v5/tree/main/onnx)
62
+
63
+ ## 使用说明
64
+
65
+ 将本仓库所有文件下载到设备中
66
+
67
+ ```
68
+ PPOCR_v6# tree -L 1
69
+ .
70
+ |-- 11.jpg # 测试图
71
+ |-- README-zh.md
72
+ |-- README.md
73
+ |-- axmodel # 各个版本的axmodel
74
+ |-- cls.json # 方向分类模型转换axmodel配置文件
75
+ |-- dataset # 量化数据集 与 测试数据集
76
+ |-- det.json # 检测模型转换axmodel配置文件
77
+ |-- fonts # 绘制字体
78
+ |-- onnx # 各模型的原始onnx文件
79
+ |-- ppocrv6_ax.py # axmodel运行pipeline
80
+ |-- ppocrv6_onnx.py # onnx运行pipeline
81
+ |-- rec.json # 识别模型转换axmodel配置文件
82
+ |-- res-ax.jpg # axmodel运行结果
83
+ |-- res-onnx.jpg # onnx运行结果
84
+ |-- run_det_ax.py # 检测axmodel模型准确率测试文件
85
+ |-- run_det_onnx.py # 检测onnx模型准确率测试文件
86
+ |-- run_rec_ax.py # 识别axmodel模型准确率测试文件
87
+ `-- run_rec_onnx.py # 识别onnx模型准确率测试文件
88
+ ```
89
+
90
+ ### 转换
91
+ ```
92
+ cd dataset
93
+ sh download_quant_dataset.sh
94
+ sh download_val_dataset.sh
95
+ cd ..
96
+ pulsar2 build --config det.json
97
+ pulsar2 build --config cls.json
98
+ pulsar2 build --config rec.json
99
+
100
+ ```
101
+
102
+ ### 测试
103
+
104
+ #### 检测
105
+
106
+ ``` python
107
+ python3 run_det_onnx.py --resize_mode letterbox # resize_mode可选:letterbox(默认)、stretch(官方)
108
+ """
109
+ Images: 50
110
+ GT boxes: 201
111
+ DET boxes: 151
112
+ Matched: 77
113
+ Precision: 0.5099 (50.99%)
114
+ Recall: 0.3831 (38.31%)
115
+ Hmean (F1): 0.4375
116
+ """
117
+ python3 run_det_ax.py --resize_mode letterbox # resize_mode可选:letterbox(默认)、stretch(官方)
118
+ """
119
+ Images: 50
120
+ GT boxes: 201
121
+ DET boxes: 150
122
+ Matched: 75
123
+ Precision: 0.5000 (50.00%)
124
+ Recall: 0.3731 (37.31%)
125
+ Hmean (F1): 0.4274
126
+ """
127
+ ```
128
+
129
+ 注:`resize_mode`为`stretch`时,为官方做法,直接`resize`到模型输入大小;`letterbox`为右下角填充,与动态输入的`onnx`差距较小,在本项测试中指标较优。可下载 [small-det-onnx](https://huggingface.co/PaddlePaddle/PP-OCRv6_small_det_onnx) 中动态输入 `shape` 的 `inference.onnx` 进行对比。
130
+
131
+ #### 识别
132
+
133
+ ``` python
134
+ python3 run_rec_onnx.py
135
+ """
136
+ Total samples: 2077
137
+ Correct (exact match): 1563
138
+ Accuracy: 0.7525 (75.25%)
139
+ Norm Edit Distance: 0.8947
140
+ """
141
+ python3 run_rec_ax.py
142
+ """
143
+ Total samples: 2077
144
+ Correct (exact match): 1518
145
+ Accuracy: 0.7309 (73.09%)
146
+ Norm Edit Distance: 0.8781
147
+ """
148
+ ```
149
+
150
+ ### 推理
151
+
152
+ 在 AX650 主机上推理,如 M4N-Dock(爱芯派Pro)
153
+
154
+ 输入图片:
155
+ ![input](./11.jpg)
156
+
157
+
158
+
159
+ ``` python
160
+ python3 ppocrv6_onnx.py --use_angle_cls --visualize --image 11.jpg
161
+ python3 ppocrv6_ax.py --use_angle_cls --visualize --image 11.jpg
162
+ ```
163
+
164
+ 输出图片:
165
+ ![output](./res-ax.jpg)
166
+
167
+ ### 其他
168
+ `tiny`的识别模型量化误差较大,指标如下:
169
+ ```
170
+ onnx-preds:
171
+ Total samples: 2077
172
+ Correct (exact match): 1271
173
+ Accuracy: 0.6119 (61.19%)
174
+ Norm Edit Distance: 0.8263
175
+
176
+ ax-w8a16-preds:
177
+ Total samples: 2077
178
+ Correct (exact match): 1178
179
+ Accuracy: 0.5672 (56.72%)
180
+ Norm Edit Distance: 0.7941
181
+ ```
182
+
183
+ #### TODO
184
+
185
+ - ax620性能测试
186
+ - ax615性能测试
README.md CHANGED
@@ -1,3 +1,185 @@
1
  ---
2
- license: apache-2.0
 
 
 
 
 
 
 
 
 
 
 
 
3
  ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ license: mit
3
+ language:
4
+ - en
5
+ base_model:
6
+ - PP-OCRv6_mobile_det
7
+ - PP-OCRv6_mobile_rec
8
+ - PP-LCNet_x0_25_textline_ori
9
+ pipeline_tag: text-classification
10
+ tags:
11
+ - OCR
12
+ - paddle
13
+ - PPOCRv6
14
+ - axera
15
  ---
16
+
17
+ # PPOCR_v6
18
+ > English | [中文](./README-zh.md)
19
+
20
+ This version of PPOCR_v6 has been converted to run on AXERA NPU with **w8a16** quantization.
21
+
22
+ ## Conversion Tool Links
23
+
24
+ If you are interested in model conversion, you can export axmodel through the following links:
25
+
26
+ - [ax-samples-github](https://github.com/AXERA-TECH/ax-samples), other interesting samples
27
+
28
+ - [Pulsar2 Documentation, ONNX to axmodel conversion](https://pulsar2-docs.readthedocs.io/en/latest/pulsar2/introduction.html)
29
+
30
+ ## Supported Platforms
31
+
32
+ - AX650
33
+ - [M4N-Dock (AXERA Pi Pro)](https://wiki.sipeed.com/hardware/zh/maixIV/m4ndock/m4ndock.html)
34
+ - [M.2 Accelerator Card](https://axcl-docs.readthedocs.io/zh-cn/latest/doc_guide_hardware.html)
35
+ - AX630C
36
+ - [AXERA Pi 2](https://axera-pi-2-docs-cn.readthedocs.io/zh-cn/latest/index.html)
37
+ - [Module-LLM](https://docs.m5stack.com/zh_CN/module/Module-LLM)
38
+ - [LLM630 Compute Kit](https://docs.m5stack.com/zh_CN/core/LLM630%20Compute%20Kit)
39
+ - AX615
40
+ - [AX615 IPC SoC](https://www.axera-tech.com/zh-hans/product/2956.html)
41
+
42
+ ### Performance Benchmarks
43
+
44
+ | Chip | Model | npu_mode | Latency (ms) |
45
+ | ----- | ------------------------------- | -------- | ------------ |
46
+ | | PP-OCRv6_small_det | NPU1 | 42.158 |
47
+ | | PP-OCRv6_small_det | NPU3 | 23.837 |
48
+ | AX650 | PP-LCNet_x0_25_textline_ori | NPU1 | 0.294 |
49
+ | | PP-LCNet_x0_25_textline_ori | NPU3 | 0.172 |
50
+ | | PP-OCRv6_small_rec | NPU1 | 2.473 |
51
+ | | PP-OCRv6_small_rec | NPU3 | 1.073 |
52
+
53
+ Benchmarking command:
54
+
55
+ ``` bash
56
+ ax_run_model -w 10 -r 100 -m xx.axmodel
57
+ ```
58
+
59
+ Recognition and detection ONNX model sources: [small-det-onnx](https://huggingface.co/PaddlePaddle/PP-OCRv6_small_det_onnx) and [small-rec-onnx](https://huggingface.co/PaddlePaddle/PP-OCRv6_small_rec_onnx) from [PaddlePaddle/PP-OCRv6](https://huggingface.co/collections/PaddlePaddle/pp-ocrv6)
60
+
61
+ Text direction classifier: [AXERA-TECH/PPOCR_v5](https://huggingface.co/AXERA-TECH/PPOCR_v5/tree/main/onnx)
62
+
63
+ ## Usage
64
+
65
+ Download all files from this repository to your device.
66
+
67
+ ```
68
+ PPOCR_v6# tree -L 1
69
+ .
70
+ |-- 11.jpg # Test image
71
+ |-- README-zh.md
72
+ |-- README.md
73
+ |-- axmodel # Axmodel files for each version
74
+ |-- cls.json # Text direction classifier axmodel conversion config
75
+ |-- dataset # Quantization dataset & test dataset
76
+ |-- det.json # Detection model axmodel conversion config
77
+ |-- fonts # Rendering fonts
78
+ |-- onnx # Original ONNX files for each model
79
+ |-- ppocrv6_ax.py # Axmodel inference pipeline
80
+ |-- ppocrv6_onnx.py # ONNX inference pipeline
81
+ |-- rec.json # Recognition model axmodel conversion config
82
+ |-- res-ax.jpg # Axmodel inference result
83
+ |-- res-onnx.jpg # ONNX inference result
84
+ |-- run_det_ax.py # Detection axmodel accuracy test script
85
+ |-- run_det_onnx.py # Detection ONNX accuracy test script
86
+ |-- run_rec_ax.py # Recognition axmodel accuracy test script
87
+ `-- run_rec_onnx.py # Recognition ONNX accuracy test script
88
+ ```
89
+
90
+ ### Conversion
91
+ ```
92
+ cd dataset
93
+ sh download_quant_dataset.sh
94
+ sh download_val_dataset.sh
95
+ cd ..
96
+ pulsar2 build --config det.json
97
+ pulsar2 build --config cls.json
98
+ pulsar2 build --config rec.json
99
+
100
+ ```
101
+
102
+ ### Testing
103
+
104
+ #### Detection
105
+
106
+ ``` python
107
+ python3 run_det_onnx.py --resize_mode letterbox # resize_mode options: letterbox (default), stretch (official)
108
+ """
109
+ Images: 50
110
+ GT boxes: 201
111
+ DET boxes: 151
112
+ Matched: 77
113
+ Precision: 0.5099 (50.99%)
114
+ Recall: 0.3831 (38.31%)
115
+ Hmean (F1): 0.4375
116
+ """
117
+ python3 run_det_ax.py --resize_mode letterbox # resize_mode options: letterbox (default), stretch (official)
118
+ """
119
+ Images: 50
120
+ GT boxes: 201
121
+ DET boxes: 150
122
+ Matched: 75
123
+ Precision: 0.5000 (50.00%)
124
+ Recall: 0.3731 (37.31%)
125
+ Hmean (F1): 0.4274
126
+ """
127
+ ```
128
+
129
+ Note: When `resize_mode` is `stretch`, it follows the official approach of directly resizing to the model input size. When `resize_mode` is `letterbox`, it pads the bottom-right corner, which has a smaller gap from the dynamic-input ONNX model and achieves better metrics in this test. You can download the `inference.onnx` with dynamic input `shape` from [small-det-onnx](https://huggingface.co/PaddlePaddle/PP-OCRv6_small_det_onnx) for comparison.
130
+
131
+ #### Recognition
132
+
133
+ ``` python
134
+ python3 run_rec_onnx.py
135
+ """
136
+ Total samples: 2077
137
+ Correct (exact match): 1563
138
+ Accuracy: 0.7525 (75.25%)
139
+ Norm Edit Distance: 0.8947
140
+ """
141
+ python3 run_rec_ax.py
142
+ """
143
+ Total samples: 2077
144
+ Correct (exact match): 1518
145
+ Accuracy: 0.7309 (73.09%)
146
+ Norm Edit Distance: 0.8781
147
+ """
148
+ ```
149
+
150
+ ### Inference
151
+
152
+ Run inference on AX650 host, such as M4N-Dock (AXERA Pi Pro).
153
+
154
+ Input image:
155
+ ![input](./11.jpg)
156
+
157
+
158
+ ``` python
159
+ python3 ppocrv6_onnx.py --use_angle_cls --visualize --image 11.jpg
160
+ python3 ppocrv6_ax.py --use_angle_cls --visualize --image 11.jpg
161
+ ```
162
+
163
+ Output image:
164
+ ![output](./res-ax.jpg)
165
+
166
+ ### Others
167
+ The `tiny` recognition model has a large quantization error. Metrics are as follows:
168
+ ```
169
+ onnx-preds:
170
+ Total samples: 2077
171
+ Correct (exact match): 1271
172
+ Accuracy: 0.6119 (61.19%)
173
+ Norm Edit Distance: 0.8263
174
+
175
+ ax-w8a16-preds:
176
+ Total samples: 2077
177
+ Correct (exact match): 1178
178
+ Accuracy: 0.5672 (56.72%)
179
+ Norm Edit Distance: 0.7941
180
+ ```
181
+
182
+ #### TODO
183
+
184
+ - ax620 performance benchmark
185
+ - ax615 performance benchmark
axmodel/ax650/cls_npu1.axmodel ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9bf1940e8047b7308045248740089fb453aeb7af641bafd10508334d71121179
3
+ size 647741
axmodel/ax650/cls_npu3.axmodel ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:db63d3946a4af62533a939a2e7121f595fdac931cddb57413d4f9a3099a0e17a
3
+ size 786685
axmodel/ax650/det_npu1.axmodel ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ddcbfb8d13ebfeaeaedecfdf87778ec86c7203cc1dbf0340031ed070e145616a
3
+ size 4712586
axmodel/ax650/det_npu3.axmodel ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:07e1f5ca9369076b0a7bafeffef2456f24394cfe9beac048e6403cd3f901c885
3
+ size 4264170
axmodel/ax650/rec_npu1.axmodel ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b32efcc3840955bec45f9b44683192494c8deee9eed86dce1a29fec2e15e24fd
3
+ size 6713642
axmodel/ax650/rec_npu3.axmodel ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:722fcd568848065549ed9d6a190c12790eae7ef81e7e5ff593905c302d90432a
3
+ size 6904074
cls.json ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "input": "./onnx/cls_mobile_sim_static.onnx",
3
+ "output_dir": "./cls-650-npu1",
4
+ "output_name": "cls_npu1.axmodel",
5
+ "model_type": "ONNX",
6
+ // "target_hardware": "AX620E",
7
+ "npu_mode": "NPU1",
8
+ "onnx_opt": {
9
+ "disable_onnx_optimization": false,
10
+ "enable_onnxsim": true,
11
+ "model_check": true,
12
+ "disable_transformation_check": false
13
+ },
14
+ "quant": {
15
+ "input_configs": [
16
+ {
17
+ "tensor_name": "x",
18
+ "calibration_dataset": "./dataset/rec_quant.tar",
19
+ "calibration_format": "Image",
20
+ "calibration_size": -1,
21
+ "calibration_mean": [127.5, 127.5, 127.5],
22
+ "calibration_std": [127.5, 127.5, 127.5]
23
+ }
24
+ ],
25
+ "calibration_method": "MinMax",
26
+ "precision_analysis": true,
27
+ "precision_analysis_method": "EndToEnd",
28
+ "precision_analysis_mode": "Reference",
29
+ "layer_configs": [
30
+ {
31
+ "start_tensor_names": ["DEFAULT"],
32
+ "end_tensor_names": ["DEFAULT"],
33
+ "data_type": "U16"
34
+ }
35
+ ]
36
+ },
37
+ "input_processors": [
38
+ {
39
+ "tensor_name": "x",
40
+ "tensor_format": "BGR",
41
+ "tensor_layout": "NCHW",
42
+ "src_format": "BGR",
43
+ "src_dtype": "FP32",
44
+ "src_layout": "NCHW",
45
+ "csc_mode": "NoCSC"
46
+ }
47
+ ],
48
+ "compiler": {
49
+ "check": 2
50
+ }
51
+ }
dataset/download_quant_dataset.sh ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ wget https://github.com/AXERA-TECH/PPOCR_v5.axera/releases/download/v1.0.1/rec_pic.tar.gz
2
+ wget https://github.com/AXERA-TECH/PPOCR_v5.axera/releases/download/v1.0.1/ocr_pic.tar.gz
dataset/download_val_dataset.sh ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ wget https://paddle-model-ecology.bj.bcebos.com/paddlex/data/ocr_det_dataset_examples.tar && tar -xvf ocr_det_dataset_examples.tar
2
+ wget https://paddle-model-ecology.bj.bcebos.com/paddlex/data/ocr_rec_dataset_examples.tar && tar -xvf ocr_rec_dataset_examples.tar
det.json ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "input": "onnx/det_inference_static_sim.onnx",
3
+ "output_dir": "./det-ax650-npu3",
4
+ "output_name": "det_npu3.axmodel",
5
+ "model_type": "ONNX",
6
+ // "target_hardware": "AX620E",
7
+ "npu_mode": "NPU3",
8
+
9
+ "quant": {
10
+ "input_configs": [
11
+ {
12
+ "tensor_name": "DEFAULT",
13
+ "calibration_dataset": "./dataset/ocr_pic.tar.gz",
14
+ "calibration_size": -1,
15
+ "calibration_mean": [123.675,116.28,103.53],
16
+ "calibration_std": [58.395,57.12,57.375]
17
+ }
18
+ ],
19
+ "calibration_method": "MinMax",
20
+ "precision_analysis": true,
21
+ "precision_analysis_method": "EndToEnd",
22
+ "precision_analysis_mode": "Reference",
23
+ "conv_bias_data_type": "FP32",
24
+ "layer_configs": [
25
+ {
26
+ "start_tensor_names": ["DEFAULT"],
27
+ "end_tensor_names": ["DEFAULT"],
28
+ "data_type": "U16"
29
+ }
30
+ ]
31
+ },
32
+ "input_processors": [
33
+ {
34
+ "tensor_name": "DEFAULT",
35
+ "tensor_format": "BGR",
36
+ "tensor_layout": "NCHW",
37
+ "src_format": "BGR",
38
+ "src_dtype": "FP32",
39
+ "src_layout": "NCHW",
40
+ "csc_mode": "NoCSC"
41
+ }
42
+ ],
43
+ "compiler": {
44
+ "check": 2
45
+ }
46
+ }
fonts/simfang.ttf ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:521c6f7546b4eb64fa4b0cd604bbd36333a20a57e388c8e2ad2ad07b9e593864
3
+ size 10576012
onnx/cls_mobile_sim_static.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:870f2ed9c720a3e955e570c5309f88ed3a8b06c02a537762a10fac33070384a2
3
+ size 980006
onnx/det_inference_static_sim.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:47b0141861a709e98871cf407924e6012cec5cf076e5073a94dd4a910039b3bd
3
+ size 9855211
onnx/rec_inference.yml ADDED
The diff for this file is too large to render. See raw diff
 
onnx/rec_inference_static_sim.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ffda730a074cbd034878d6ad2d47915d1307cd3a007bcc606a5763b89e5637f3
3
+ size 21106528
ppocrv6_ax.py ADDED
@@ -0,0 +1,721 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ PP-OCRv6 ONNX Inference (standalone, zero Paddle dependency)
4
+
5
+ Dependencies:
6
+ numpy, opencv-python, onnxruntime, pyyaml, shapely, pyclipper
7
+
8
+ Usage:
9
+ from ppocrv6_onnx import PPOCRv6Onnx
10
+ ocr = PPOCRv6Onnx(det_onnx="det.onnx", rec_onnx="rec.onnx", char_dict="inference.yml")
11
+ results = ocr(img) # img is BGR numpy array
12
+
13
+ # With direction classifier:
14
+ ocr = PPOCRv6Onnx(..., cls_onnx="cls.onnx", cls_label_list=["0","180"], cls_thresh=0.9)
15
+ """
16
+
17
+ import argparse
18
+ import json
19
+ import math
20
+ import os
21
+ from typing import List, Optional, Tuple, Union
22
+
23
+ import cv2
24
+ import numpy as np
25
+ import axengine as ort
26
+ import yaml
27
+
28
+ from PIL import Image, ImageDraw, ImageFont
29
+ from shapely.geometry import Polygon
30
+ import pyclipper
31
+ import random
32
+
33
+
34
+ # ============================================================================
35
+ # Helpers
36
+ # ============================================================================
37
+
38
+ def _get_dim_value(dim):
39
+ """Extract integer value from an ONNX Runtime dimension, returning 0 for dynamic."""
40
+ if dim is None:
41
+ return 0
42
+ if isinstance(dim, str):
43
+ return 0
44
+ if hasattr(dim, "dim_value"):
45
+ return int(dim.dim_value) if dim.dim_value else 0
46
+ if hasattr(dim, "dim_param"):
47
+ return 0
48
+ try:
49
+ return int(dim)
50
+ except (TypeError, ValueError):
51
+ return 0
52
+
53
+
54
+ def _detect_fixed_dims(session: ort.InferenceSession, det_onnx: str):
55
+ """Detect fixed H/W from ONNX input shape. Returns (fixed_h, fixed_w)."""
56
+ inp = session.get_inputs()[0]
57
+ h, w = _get_dim_value(inp.shape[2]), _get_dim_value(inp.shape[3])
58
+ if h == 0 and w == 0:
59
+ try:
60
+ import onnx
61
+ m = onnx.load(det_onnx)
62
+ dims = m.graph.input[0].type.tensor_type.shape.dim
63
+ h = dims[2].dim_value if len(dims) > 2 else 0
64
+ w = dims[3].dim_value if len(dims) > 3 else 0
65
+ except Exception:
66
+ pass
67
+ return (h if h > 0 else 0), (w if w > 0 else 0)
68
+
69
+
70
+ def _load_char_dict(source: Union[str, List[str]]) -> List[str]:
71
+ if isinstance(source, list):
72
+ return source
73
+ ext = os.path.splitext(source)[1].lower()
74
+ if ext in (".yml", ".yaml"):
75
+ with open(source, "r", encoding="utf-8") as f:
76
+ cfg = yaml.safe_load(f)
77
+ dic = cfg.get("PostProcess", {}).get("character_dict", [])
78
+ if not dic:
79
+ raise ValueError(f"No PostProcess.character_dict found in {source}")
80
+ return dic
81
+ elif ext == ".txt":
82
+ with open(source, "r", encoding="utf-8") as f:
83
+ return [line.strip("\n\r") for line in f.readlines()]
84
+ else:
85
+ raise ValueError(f"Unsupported char_dict source: {source}. Use .yml, .txt, or list.")
86
+
87
+
88
+ # ============================================================================
89
+ # 1. Detection Preprocessing
90
+ # ============================================================================
91
+
92
+ class _DetResizeForTest:
93
+ def __init__(self, limit_side_len=960, limit_type="max", max_side_limit=4000):
94
+ self.limit_side_len = limit_side_len
95
+ self.limit_type = limit_type
96
+ self.max_side_limit = max_side_limit
97
+
98
+ def _image_padding(self, im, value=0):
99
+ h, w, c = im.shape
100
+ im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value
101
+ im_pad[:h, :w, :] = im
102
+ return im_pad
103
+
104
+ def _resize_image_type0(self, img):
105
+ h, w, _ = img.shape
106
+ limit_side_len = self.limit_side_len
107
+ if self.limit_type == "max":
108
+ ratio = float(limit_side_len) / max(h, w) if max(h, w) > limit_side_len else 1.0
109
+ elif self.limit_type == "min":
110
+ ratio = float(limit_side_len) / min(h, w) if min(h, w) < limit_side_len else 1.0
111
+ elif self.limit_type == "resize_long":
112
+ ratio = float(limit_side_len) / max(h, w)
113
+ else:
114
+ raise ValueError(f"not support limit_type: {self.limit_type}")
115
+
116
+ resize_h, resize_w = int(h * ratio), int(w * ratio)
117
+ if max(resize_h, resize_w) > self.max_side_limit:
118
+ ratio = float(self.max_side_limit) / max(resize_h, resize_w)
119
+ resize_h, resize_w = int(resize_h * ratio), int(resize_w * ratio)
120
+ resize_h = max(int(round(resize_h / 32) * 32), 32)
121
+ resize_w = max(int(round(resize_w / 32) * 32), 32)
122
+ if int(resize_w) <= 0 or int(resize_h) <= 0:
123
+ return None, (None, None)
124
+ img = cv2.resize(img, (int(resize_w), int(resize_h)))
125
+ ratio_h, ratio_w = resize_h / float(h), resize_w / float(w)
126
+ return img, [ratio_h, ratio_w]
127
+
128
+ def __call__(self, img):
129
+ src_h, src_w = img.shape[:2]
130
+ if sum([src_h, src_w]) < 64:
131
+ img = self._image_padding(img)
132
+ img, [ratio_h, ratio_w] = self._resize_image_type0(img)
133
+ shape = np.array([src_h, src_w, ratio_h, ratio_w])
134
+ return img, shape
135
+
136
+
137
+ class _NormalizeImage:
138
+ def __init__(self, mean, std, scale=1.0 / 255.0, order="hwc"):
139
+ self.scale = np.float32(scale)
140
+ shape = (1, 1, 3) if order == "hwc" else (3, 1, 1)
141
+ self.mean = np.array(mean, dtype=np.float32).reshape(shape)
142
+ self.std = np.array(std, dtype=np.float32).reshape(shape)
143
+
144
+ def __call__(self, img):
145
+ return (img.astype("float32") * self.scale - self.mean) / self.std
146
+
147
+
148
+ class _ToCHWImage:
149
+ def __call__(self, img):
150
+ return img.transpose((2, 0, 1))
151
+
152
+
153
+ # ============================================================================
154
+ # 2. Recognition Preprocessing
155
+ # ============================================================================
156
+
157
+ def _resize_norm_img(img, image_shape, max_wh_ratio=None):
158
+ imgC, imgH, imgW = image_shape
159
+ if max_wh_ratio is None:
160
+ max_wh_ratio = imgW * 1.0 / imgH
161
+ h, w = img.shape[:2]
162
+ max_wh_ratio = max(max_wh_ratio, w / h)
163
+ target_w = int(imgH * max_wh_ratio)
164
+ h, w = img.shape[:2]
165
+ ratio = w / h
166
+ resized_w = target_w if math.ceil(imgH * ratio) > target_w else int(math.ceil(imgH * ratio))
167
+ resized_image = cv2.resize(img, (resized_w, imgH)).astype("float32")
168
+ resized_image = resized_image.transpose((2, 0, 1))
169
+ # resized_image /= 255.0
170
+ # resized_image -= 0.5
171
+ # resized_image /= 0.5
172
+ padding_im = np.zeros((imgC, imgH, target_w), dtype=np.float32)
173
+ padding_im[:, :, 0:resized_w] = resized_image
174
+ return padding_im
175
+
176
+
177
+ # ============================================================================
178
+ # 3. Detection Postprocessing
179
+ # ============================================================================
180
+
181
+ class _DBPostProcess:
182
+ def __init__(self, thresh=0.3, box_thresh=0.7, max_candidates=1000,
183
+ unclip_ratio=2.0, use_dilation=False, score_mode="fast", box_type="quad"):
184
+ self.thresh = thresh
185
+ self.box_thresh = box_thresh
186
+ self.max_candidates = max_candidates
187
+ self.unclip_ratio = unclip_ratio
188
+ self.min_size = 3
189
+ self.score_mode = score_mode
190
+ self.box_type = box_type
191
+ self.dilation_kernel = None if not use_dilation else np.array([[1, 1], [1, 1]])
192
+
193
+ def _unclip(self, box, unclip_ratio):
194
+ poly = Polygon(box)
195
+ distance = poly.area * unclip_ratio / poly.length
196
+ offset = pyclipper.PyclipperOffset()
197
+ offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
198
+ return offset.Execute(distance)
199
+
200
+ def _get_mini_boxes(self, contour):
201
+ bb = cv2.minAreaRect(contour)
202
+ points = sorted(list(cv2.boxPoints(bb)), key=lambda x: x[0])
203
+ i1, i4 = (0, 1) if points[1][1] > points[0][1] else (1, 0)
204
+ i2, i3 = (2, 3) if points[3][1] > points[2][1] else (3, 2)
205
+ return [points[i1], points[i2], points[i3], points[i4]], min(bb[1])
206
+
207
+ def _box_score_fast(self, bitmap, _box):
208
+ h, w = bitmap.shape[:2]
209
+ box = _box.copy()
210
+ xmin = np.clip(np.floor(box[:, 0].min()).astype("int32"), 0, w - 1)
211
+ xmax = np.clip(np.ceil(box[:, 0].max()).astype("int32"), 0, w - 1)
212
+ ymin = np.clip(np.floor(box[:, 1].min()).astype("int32"), 0, h - 1)
213
+ ymax = np.clip(np.ceil(box[:, 1].max()).astype("int32"), 0, h - 1)
214
+ mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
215
+ box[:, 0] -= xmin
216
+ box[:, 1] -= ymin
217
+ cv2.fillPoly(mask, box.reshape(1, -1, 2).astype("int32"), 1)
218
+ return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
219
+
220
+ def _boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
221
+ height, width = _bitmap.shape
222
+ outs = cv2.findContours(
223
+ (_bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
224
+ )
225
+ contours = outs[0] if len(outs) == 2 else outs[1]
226
+ num = min(len(contours), self.max_candidates)
227
+ boxes, scores = [], []
228
+ for i in range(num):
229
+ points, sside = self._get_mini_boxes(contours[i])
230
+ if sside < self.min_size:
231
+ continue
232
+ points = np.array(points)
233
+ score = self._box_score_fast(pred, points.reshape(-1, 2))
234
+ if self.box_thresh > score:
235
+ continue
236
+ box = self._unclip(points, self.unclip_ratio)
237
+ if len(box) > 1:
238
+ continue
239
+ box = np.array(box).reshape(-1, 1, 2)
240
+ box, sside = self._get_mini_boxes(box)
241
+ if sside < self.min_size + 2:
242
+ continue
243
+ box = np.array(box)
244
+ box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
245
+ box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
246
+ boxes.append(box.astype("int32"))
247
+ scores.append(score)
248
+ return np.array(boxes, dtype="int32"), scores
249
+
250
+ def __call__(self, pred, shape_list):
251
+ pred = pred[:, 0, :, :]
252
+ segmentation = pred > self.thresh
253
+ boxes_batch = []
254
+ for bi in range(pred.shape[0]):
255
+ src_h, src_w, ratio_h, ratio_w = shape_list[bi]
256
+ mask = cv2.dilate(np.array(segmentation[bi]).astype(np.uint8),
257
+ self.dilation_kernel) if self.dilation_kernel is not None else segmentation[bi]
258
+ boxes, _ = self._boxes_from_bitmap(pred[bi], mask, src_w, src_h)
259
+ boxes_batch.append(boxes)
260
+ return boxes_batch
261
+
262
+
263
+ # ============================================================================
264
+ # 4. Recognition Postprocessing
265
+ # ============================================================================
266
+
267
+ class _CTCLabelDecode:
268
+ def __init__(self, character_list: List[str], use_space_char=True):
269
+ self.character_str = list(character_list)
270
+ if use_space_char:
271
+ self.character_str.append(" ")
272
+ dict_character = ["blank"] + self.character_str
273
+ self.character = dict_character
274
+
275
+ def decode(self, text_index, text_prob=None, is_remove_duplicate=True):
276
+ result_list = []
277
+ for bi in range(len(text_index)):
278
+ sel = np.ones(len(text_index[bi]), dtype=bool)
279
+ if is_remove_duplicate:
280
+ sel[1:] = text_index[bi][1:] != text_index[bi][:-1]
281
+ sel &= text_index[bi] != 0
282
+ chars = [self.character[int(t)] for t in text_index[bi][sel]]
283
+ conf = text_prob[bi][sel] if text_prob is not None else [1] * len(sel)
284
+ if len(conf) == 0:
285
+ conf = [0]
286
+ result_list.append(("".join(chars), float(np.mean(conf))))
287
+ return result_list
288
+
289
+ def __call__(self, preds):
290
+ return self.decode(preds.argmax(axis=2), preds.max(axis=2), is_remove_duplicate=True)
291
+
292
+
293
+ # ============================================================================
294
+ # 5. Image Utilities
295
+ # ============================================================================
296
+
297
+ def _get_rotate_crop_image(img: np.ndarray, points: np.ndarray) -> np.ndarray:
298
+ assert len(points) == 4
299
+ cw = int(max(np.linalg.norm(points[0] - points[1]), np.linalg.norm(points[2] - points[3])))
300
+ ch = int(max(np.linalg.norm(points[0] - points[3]), np.linalg.norm(points[1] - points[2])))
301
+ pts_std = np.float32([[0, 0], [cw, 0], [cw, ch], [0, ch]])
302
+ M = cv2.getPerspectiveTransform(points.astype(np.float32), pts_std)
303
+ dst = cv2.warpPerspective(img, M, (cw, ch), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC)
304
+ if dst.shape[0] * 1.0 / dst.shape[1] >= 1.5:
305
+ dst = np.rot90(dst)
306
+ return dst
307
+
308
+
309
+ def _sorted_boxes(dt_boxes):
310
+ if len(dt_boxes) == 0:
311
+ return dt_boxes
312
+ boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
313
+ lst = list(boxes)
314
+ for i in range(len(lst) - 1):
315
+ for j in range(i, -1, -1):
316
+ if abs(lst[j + 1][0][1] - lst[j][0][1]) < 10 and lst[j + 1][0][0] < lst[j][0][0]:
317
+ lst[j], lst[j + 1] = lst[j + 1], lst[j]
318
+ else:
319
+ break
320
+ return lst
321
+
322
+
323
+ def draw_ocr_result(
324
+ img: np.ndarray,
325
+ results: List[dict],
326
+ font_path: str = "./fonts/simfang.ttf",
327
+ ) -> np.ndarray:
328
+ """Draw detection boxes (semi-transparent) on original image, with text list on the right side."""
329
+ h, w = img.shape[:2]
330
+
331
+ # --- left: original image with semi-transparent colored boxes ---
332
+ pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
333
+ overlay = Image.new("RGBA", pil_img.size, (0, 0, 0, 0))
334
+ draw_overlay = ImageDraw.Draw(overlay)
335
+
336
+ random.seed(0)
337
+ for res in results:
338
+ box = res["box"]
339
+ color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 90)
340
+ draw_overlay.polygon([tuple(p) for p in box], fill=color)
341
+
342
+ left_img = Image.alpha_composite(pil_img.convert("RGBA"), overlay).convert("RGB")
343
+
344
+ # --- right: white canvas with text list ---
345
+ right_w = int(w * 0.9)
346
+ right = Image.new("RGB", (right_w, h), (255, 255, 255))
347
+ draw_right = ImageDraw.Draw(right)
348
+
349
+ try:
350
+ font = ImageFont.truetype(font_path, 14)
351
+ except (OSError, IOError):
352
+ font = ImageFont.load_default()
353
+
354
+ y = 5
355
+ gap = 18
356
+ for i, res in enumerate(results):
357
+ text = f"{i+1}. {res['text']} ({res['confidence']:.3f})"
358
+ # Color block matching the box
359
+ random.seed(i)
360
+ blk_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
361
+ draw_right.rectangle([5, y + 3, 15, y + 14], fill=blk_color, outline=(0, 0, 0))
362
+ draw_right.text((20, y), text, fill=(0, 0, 0), font=font)
363
+ y += gap
364
+
365
+ # --- concat left + right ---
366
+ result_img = Image.new("RGB", (w + right_w, h))
367
+ result_img.paste(left_img, (0, 0))
368
+ result_img.paste(right, (w, 0))
369
+
370
+ return cv2.cvtColor(np.array(result_img), cv2.COLOR_RGB2BGR)
371
+
372
+
373
+ # ============================================================================
374
+ # 6. Main Inference Engine
375
+ # ============================================================================
376
+
377
+ class PPOCRv6Onnx:
378
+
379
+ def __init__(
380
+ self,
381
+ det_onnx: str,
382
+ rec_onnx: str,
383
+ char_dict: Union[str, List[str]],
384
+ # Detection params
385
+ det_limit_side_len: int = 960,
386
+ det_db_thresh: float = 0.2,
387
+ det_db_box_thresh: float = 0.4,
388
+ det_db_unclip_ratio: float = 1.4,
389
+ det_max_candidates: int = 3000,
390
+ # Recognition params
391
+ rec_image_shape: Tuple[int, int, int] = (3, 48, 320),
392
+ rec_batch_num: int = 1,
393
+ # Classifier params
394
+ use_angle_cls: bool = False,
395
+ cls_onnx: Optional[str] = None,
396
+ cls_image_shape: Tuple[int, int, int] = (3, 48, 192),
397
+ cls_batch_num: int = 1,
398
+ cls_thresh: float = 0.9,
399
+ cls_label_list: Optional[List[str]] = None,
400
+ # Common
401
+ drop_score: float = 0.5,
402
+ use_gpu: bool = False,
403
+ onnx_providers: Optional[List[str]] = None,
404
+ resize_mode: str = "letterbox",
405
+ ):
406
+ assert resize_mode in ("letterbox", "stretch"), f"invalid resize_mode: {resize_mode}"
407
+ if cls_label_list is None:
408
+ cls_label_list = ["0", "180"]
409
+
410
+ self.rec_image_shape = rec_image_shape
411
+ self.rec_batch_num = rec_batch_num
412
+ self.drop_score = drop_score
413
+ self.use_angle_cls = use_angle_cls
414
+ self.cls_thresh = cls_thresh
415
+ self.cls_label_list = cls_label_list
416
+ self.cls_batch_num = cls_batch_num
417
+ self.cls_image_shape = cls_image_shape
418
+ self._resize_mode = resize_mode
419
+
420
+ self.det_session = ort.InferenceSession(det_onnx)
421
+ self.det_input_name = self.det_session.get_inputs()[0].name
422
+
423
+ self.rec_session = ort.InferenceSession(rec_onnx)
424
+ self.rec_input_name = self.rec_session.get_inputs()[0].name
425
+
426
+ # Classifier session
427
+ if use_angle_cls:
428
+ if cls_onnx is None:
429
+ raise ValueError("cls_onnx is required when use_angle_cls=True")
430
+ self.cls_session = ort.InferenceSession(cls_onnx)
431
+ self.cls_input_name = self.cls_session.get_inputs()[0].name
432
+ # Detect fixed dims for cls
433
+ cls_h, cls_w = _detect_fixed_dims(self.cls_session, cls_onnx)
434
+ self._cls_fixed_h = cls_h if cls_h > 0 else 0
435
+ self._cls_fixed_w = cls_w if cls_w > 0 else 0
436
+ else:
437
+ self.cls_session = None
438
+
439
+ # Detect fixed dims for det/rec
440
+ self._det_fixed_h, self._det_fixed_w = _detect_fixed_dims(self.det_session, det_onnx)
441
+ det_shape = self.det_session.get_inputs()[0].shape
442
+ print(f"[PPOCRv6] det shape={det_shape}, fixed_h={self._det_fixed_h}, fixed_w={self._det_fixed_w}, "
443
+ f"cls={use_angle_cls}, resize_mode={resize_mode}")
444
+
445
+ # Fixed rec width from ONNX
446
+ rec_inp = self.rec_session.get_inputs()[0]
447
+ rec_fw = _get_dim_value(rec_inp.shape[3])
448
+ self._rec_fixed_w = rec_fw if rec_fw > 0 else 0
449
+
450
+ # ---- Detection pre/post ----
451
+ self._det_resize = _DetResizeForTest(limit_side_len=det_limit_side_len, limit_type="max")
452
+ # self._det_normalize = _NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
453
+ self._det_normalize = _NormalizeImage(mean=[0., 0., 0.], std=[1.0, 1.0, 1.0], scale=1.0)
454
+ self._det_to_chw = _ToCHWImage()
455
+ self._det_post = _DBPostProcess(thresh=det_db_thresh, box_thresh=det_db_box_thresh,
456
+ unclip_ratio=det_db_unclip_ratio, max_candidates=det_max_candidates, box_type="quad")
457
+
458
+ # ---- Recognition post ----
459
+ self._rec_post = _CTCLabelDecode(_load_char_dict(char_dict), use_space_char=True)
460
+
461
+ # ---- Detection with static shape support ----
462
+
463
+ def _preprocess_det(self, img: np.ndarray):
464
+ src_h, src_w = img.shape[:2]
465
+ fh, fw = self._det_fixed_h, self._det_fixed_w
466
+ # print(f'fh {fh} fw {fw}')
467
+ if fh > 0 and fw > 0:
468
+ if self._resize_mode == "stretch":
469
+ # Direct resize (official behavior)
470
+ img_r = cv2.resize(img, (fw, fh))
471
+ shape = np.array([src_h, src_w, float(fh) / src_h, float(fw) / src_w])
472
+ else:
473
+ # Letterbox: ratio-preserving + pad
474
+ ratio = min(fh / src_h, fw / src_w)
475
+ new_h = min(max(int(round(src_h * ratio / 32) * 32), 32), fh)
476
+ new_w = min(max(int(round(src_w * ratio / 32) * 32), 32), fw)
477
+ img_r = cv2.resize(img, (new_w, new_h))
478
+ pad_h, pad_w = max(0, fh - new_h), max(0, fw - new_w)
479
+ if pad_h or pad_w:
480
+ img_r = cv2.copyMakeBorder(img_r, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=(0, 0, 0))
481
+ shape = np.array([src_h * fh / new_h, src_w * fw / new_w,
482
+ float(new_h) / src_h, float(new_w) / src_w])
483
+ else:
484
+ img_r, shape = self._det_resize(img)
485
+
486
+ img_n = self._det_normalize(img_r)
487
+ img_c = self._det_to_chw(img_n)
488
+ return np.expand_dims(img_c.astype(np.float32), axis=0), shape
489
+
490
+ def _postprocess_det(self, output, shape):
491
+ return self._det_post(output, np.expand_dims(shape, axis=0))[0]
492
+
493
+ def detect(self, img):
494
+ tensor, shape = self._preprocess_det(img)
495
+ out = self.det_session.run(None, {self.det_input_name: tensor})
496
+ return self._postprocess_det(out[0], shape)
497
+
498
+ # ---- Classifier ----
499
+
500
+ def _preprocess_cls(self, img_list):
501
+ num = len(img_list)
502
+ width_list = [im.shape[1] / float(im.shape[0]) for im in img_list]
503
+ indices = np.argsort(np.array(width_list))
504
+ batches, idx_maps = [], []
505
+ for beg in range(0, num, self.cls_batch_num):
506
+ end = min(num, beg + self.cls_batch_num)
507
+ imgC, imgH, imgW = self.cls_image_shape
508
+
509
+ # Override with ONNX fixed dimensions
510
+ if self._cls_fixed_h > 0:
511
+ imgH = self._cls_fixed_h
512
+ if self._cls_fixed_w > 0:
513
+ imgW = self._cls_fixed_w
514
+
515
+ max_wh_ratio = imgW / imgH
516
+ for ino in range(beg, end):
517
+ h, w = img_list[indices[ino]].shape[:2]
518
+ max_wh_ratio = max(max_wh_ratio, w / h)
519
+ if self._cls_fixed_w > 0:
520
+ max_wh_ratio = self._cls_fixed_w / imgH
521
+
522
+ shape = (imgC, imgH, imgW)
523
+ norm_list, idx_list = [], []
524
+ for ino in range(beg, end):
525
+ orig_idx = indices[ino]
526
+ norm = _resize_norm_img(img_list[orig_idx], shape, max_wh_ratio=max_wh_ratio)
527
+ norm_list.append(np.expand_dims(norm, axis=0))
528
+ idx_list.append(orig_idx)
529
+ if norm_list:
530
+ batches.append(np.concatenate(norm_list, axis=0).astype(np.float32))
531
+ idx_maps.append(idx_list)
532
+ return batches, idx_maps
533
+
534
+ def _postprocess_cls(self, batch_outputs, idx_maps, total_num, img_list):
535
+ results = [("0", 1.0)] * total_num
536
+ for preds_batch, idx_list in zip(batch_outputs, idx_maps):
537
+ pred_ids = preds_batch.argmax(axis=1)
538
+ for i, orig_idx in enumerate(idx_list):
539
+ label = self.cls_label_list[int(pred_ids[i])]
540
+ score = float(preds_batch[i, int(pred_ids[i])])
541
+ results[orig_idx] = (label, score)
542
+ if "180" in str(label) and score > self.cls_thresh:
543
+ img_list[orig_idx] = cv2.rotate(img_list[orig_idx], cv2.ROTATE_180)
544
+ return results
545
+
546
+ def classify(self, img_list):
547
+ if not img_list or not self.use_angle_cls:
548
+ return img_list, [], 0
549
+ img_list = [im.copy() for im in img_list]
550
+ batches, idx_maps = self._preprocess_cls(img_list)
551
+ outputs = []
552
+ for batch in batches:
553
+ out = self.cls_session.run(None, {self.cls_input_name: batch})
554
+ outputs.append(out[0])
555
+ cls_res = self._postprocess_cls(outputs, idx_maps, len(img_list), img_list)
556
+ return img_list, cls_res, 0
557
+
558
+ # ---- Recognition ----
559
+
560
+ def _preprocess_rec(self, img_crop_list):
561
+ num = len(img_crop_list)
562
+ width_list = [im.shape[1] / float(im.shape[0]) for im in img_crop_list]
563
+ indices = np.argsort(np.array(width_list))
564
+ batches, idx_maps = [], []
565
+ for beg in range(0, num, self.rec_batch_num):
566
+ end = min(num, beg + self.rec_batch_num)
567
+ imgC, imgH, imgW = self.rec_image_shape
568
+ max_wh_ratio = imgW / imgH
569
+ for ino in range(beg, end):
570
+ h, w = img_crop_list[indices[ino]].shape[:2]
571
+ max_wh_ratio = max(max_wh_ratio, w / h)
572
+ # Use fixed width if set by ONNX
573
+ if self._rec_fixed_w > 0:
574
+ max_wh_ratio = self._rec_fixed_w / imgH
575
+ norm_list, idx_list = [], []
576
+ for ino in range(beg, end):
577
+ orig_idx = indices[ino]
578
+ norm = _resize_norm_img(img_crop_list[orig_idx], self.rec_image_shape, max_wh_ratio=max_wh_ratio)
579
+ norm_list.append(np.expand_dims(norm, axis=0))
580
+ idx_list.append(orig_idx)
581
+ if norm_list:
582
+ batches.append(np.concatenate(norm_list, axis=0).astype(np.float32))
583
+ idx_maps.append(idx_list)
584
+ return batches, idx_maps
585
+
586
+ def _postprocess_rec(self, batch_outputs, idx_maps, total_num):
587
+ results = [("", 0.0)] * total_num
588
+ # Decode each batch separately (different T per batch)
589
+ for preds_batch, idx_list in zip(batch_outputs, idx_maps):
590
+ texts = self._rec_post(preds_batch)
591
+ for i, orig_idx in enumerate(idx_list):
592
+ results[orig_idx] = texts[i]
593
+ return results
594
+
595
+ def recognize(self, img_crop_list):
596
+ if not img_crop_list:
597
+ return []
598
+ batches, idx_maps = self._preprocess_rec(img_crop_list)
599
+ outputs = []
600
+ for batch in batches:
601
+ out = self.rec_session.run(None, {self.rec_input_name: batch})
602
+ outputs.append(out[0])
603
+ return self._postprocess_rec(outputs, idx_maps, len(img_crop_list))
604
+
605
+ # ---- Full pipeline ----
606
+
607
+ def predict_image(self, image_path, visualize=False):
608
+ img = cv2.imread(image_path)
609
+ if img is None:
610
+ raise FileNotFoundError(f"Cannot read image: {image_path}")
611
+ return self(img, visualize=visualize)
612
+
613
+ def __call__(self, img: np.ndarray, visualize=False, use_cls=None):
614
+ ori_im = img.copy()
615
+
616
+ boxes = self.detect(img)
617
+ # print(f"[PPOCRv6] Detected {len(boxes)}")
618
+ if len(boxes) == 0:
619
+ return [] if not visualize else ori_im
620
+
621
+ boxes = _sorted_boxes(boxes)
622
+ # print(f"[PPOCRv6] Detected _sorted_boxes {len(boxes)}")
623
+
624
+ img_crop_list = []
625
+ for i, box in enumerate(boxes):
626
+ crop = _get_rotate_crop_image(ori_im, np.array(box, dtype=np.float32))
627
+ img_crop_list.append(crop)
628
+
629
+ # Direction classifier
630
+ do_cls = self.use_angle_cls if use_cls is None else use_cls
631
+ if do_cls and self.cls_session is not None:
632
+ img_crop_list, cls_res, _ = self.classify(img_crop_list)
633
+
634
+ rec_res = self.recognize(img_crop_list)
635
+
636
+ results = []
637
+ for box, (text, conf) in zip(boxes, rec_res):
638
+ if conf >= self.drop_score:
639
+ results.append({"text": text, "confidence": round(conf, 4), "box": box.tolist()})
640
+
641
+ return (results, draw_ocr_result(ori_im, results)) if visualize else (results, None)
642
+
643
+
644
+ # ============================================================================
645
+ # 7. CLI
646
+ # ============================================================================
647
+
648
+ def main():
649
+ parser = argparse.ArgumentParser(description="PP-OCRv6 ONNX Inference (standalone, no Paddle dependency)")
650
+ parser.add_argument("--det_onnx", type=str, default="axmodel/ax650/det_npu1.axmodel")
651
+ parser.add_argument("--rec_onnx", type=str, default="axmodel/ax650/rec_npu1.axmodel")
652
+ parser.add_argument("--char_dict", type=str, default="onnx/rec_inference.yml")
653
+ parser.add_argument("--image", required=True, help="Input image path")
654
+ parser.add_argument("--use_gpu", action="store_true")
655
+ parser.add_argument("--drop_score", type=float, default=0.5)
656
+ parser.add_argument("--det_limit_side_len", type=int, default=960)
657
+ parser.add_argument("--det_db_thresh", type=float, default=0.2)
658
+ parser.add_argument("--det_db_box_thresh", type=float, default=0.45)
659
+ parser.add_argument("--det_db_unclip_ratio", type=float, default=1.4)
660
+ parser.add_argument("--rec_batch_num", type=int, default=1)
661
+ parser.add_argument("--resize_mode", type=str, default="letterbox", choices=["letterbox", "stretch"])
662
+ # Classifier
663
+ parser.add_argument("--use_angle_cls", action="store_true", help="Enable direction classifier")
664
+ parser.add_argument("--cls_onnx", type=str, default="axmodel/ax650/cls_npu1.axmodel", help="Classifer ONNX model path")
665
+ parser.add_argument("--cls_thresh", type=float, default=0.9, help="Angie classifier confidence threshold")
666
+ parser.add_argument("--cls_batch_num", type=int, default=1)
667
+ # Output
668
+ parser.add_argument("--visualize", action="store_true")
669
+ parser.add_argument("--output", type=str, default=None)
670
+ parser.add_argument("--json", type=str, default=None)
671
+
672
+ args = parser.parse_args()
673
+
674
+ char_dict_src = args.char_dict
675
+ if not os.path.exists(char_dict_src) and ("," in char_dict_src or char_dict_src.startswith("[")):
676
+ char_dict = [c.strip() for c in char_dict_src.strip("[]").split(",") if c.strip()]
677
+ else:
678
+ char_dict = char_dict_src
679
+
680
+ ocr = PPOCRv6Onnx(
681
+ det_onnx=args.det_onnx,
682
+ rec_onnx=args.rec_onnx,
683
+ char_dict=char_dict,
684
+ det_limit_side_len=args.det_limit_side_len,
685
+ det_db_thresh=args.det_db_thresh,
686
+ det_db_box_thresh=args.det_db_box_thresh,
687
+ det_db_unclip_ratio=args.det_db_unclip_ratio,
688
+ rec_batch_num=args.rec_batch_num,
689
+ use_angle_cls=args.use_angle_cls,
690
+ cls_onnx=args.cls_onnx,
691
+ cls_thresh=args.cls_thresh,
692
+ cls_batch_num=args.cls_batch_num,
693
+ drop_score=args.drop_score,
694
+ use_gpu=args.use_gpu,
695
+ resize_mode=args.resize_mode,
696
+ )
697
+
698
+ do_viz = args.visualize or args.output is not None
699
+ img = cv2.imread(args.image)
700
+ if img is None:
701
+ raise FileNotFoundError(f"Cannot read image: {args.image}")
702
+
703
+ if do_viz:
704
+ results, vis = ocr(img, visualize=True)
705
+ out_path = args.output or "res-ax.jpg"
706
+ cv2.imwrite(out_path, vis)
707
+ print(f"Annotated image saved to: {out_path}")
708
+ else:
709
+ results, vis = ocr(img)
710
+
711
+ if args.json:
712
+ with open(args.json, "w", encoding="utf-8") as f:
713
+ json.dump(results, f, ensure_ascii=False, indent=2)
714
+ print(f"Results saved to: {args.json}")
715
+ else:
716
+ for i, res in enumerate(results):
717
+ print(f"{i+1}. {res['text']} ({res['confidence']:.3f})")
718
+
719
+
720
+ if __name__ == "__main__":
721
+ main()
ppocrv6_onnx.py ADDED
@@ -0,0 +1,721 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ PP-OCRv6 ONNX Inference (standalone, zero Paddle dependency)
4
+
5
+ Dependencies:
6
+ numpy, opencv-python, onnxruntime, pyyaml, shapely, pyclipper
7
+
8
+ Usage:
9
+ from ppocrv6_onnx import PPOCRv6Onnx
10
+ ocr = PPOCRv6Onnx(det_onnx="det.onnx", rec_onnx="rec.onnx", char_dict="inference.yml")
11
+ results = ocr(img) # img is BGR numpy array
12
+
13
+ # With direction classifier:
14
+ ocr = PPOCRv6Onnx(..., cls_onnx="cls.onnx", cls_label_list=["0","180"], cls_thresh=0.9)
15
+ """
16
+
17
+ import argparse
18
+ import json
19
+ import math
20
+ import os
21
+ from typing import List, Optional, Tuple, Union
22
+
23
+ import cv2
24
+ import numpy as np
25
+ import onnxruntime as ort
26
+ import yaml
27
+
28
+ from PIL import Image, ImageDraw, ImageFont
29
+ from shapely.geometry import Polygon
30
+ import pyclipper
31
+ import random
32
+
33
+
34
+ # ============================================================================
35
+ # Helpers
36
+ # ============================================================================
37
+
38
+ def _get_dim_value(dim):
39
+ """Extract integer value from an ONNX Runtime dimension, returning 0 for dynamic."""
40
+ if dim is None:
41
+ return 0
42
+ if isinstance(dim, str):
43
+ return 0
44
+ if hasattr(dim, "dim_value"):
45
+ return int(dim.dim_value) if dim.dim_value else 0
46
+ if hasattr(dim, "dim_param"):
47
+ return 0
48
+ try:
49
+ return int(dim)
50
+ except (TypeError, ValueError):
51
+ return 0
52
+
53
+
54
+ def _detect_fixed_dims(session: ort.InferenceSession, det_onnx: str):
55
+ """Detect fixed H/W from ONNX input shape. Returns (fixed_h, fixed_w)."""
56
+ inp = session.get_inputs()[0]
57
+ h, w = _get_dim_value(inp.shape[2]), _get_dim_value(inp.shape[3])
58
+ if h == 0 and w == 0:
59
+ try:
60
+ import onnx
61
+ m = onnx.load(det_onnx)
62
+ dims = m.graph.input[0].type.tensor_type.shape.dim
63
+ h = dims[2].dim_value if len(dims) > 2 else 0
64
+ w = dims[3].dim_value if len(dims) > 3 else 0
65
+ except Exception:
66
+ pass
67
+ return (h if h > 0 else 0), (w if w > 0 else 0)
68
+
69
+
70
+ def _load_char_dict(source: Union[str, List[str]]) -> List[str]:
71
+ if isinstance(source, list):
72
+ return source
73
+ ext = os.path.splitext(source)[1].lower()
74
+ if ext in (".yml", ".yaml"):
75
+ with open(source, "r", encoding="utf-8") as f:
76
+ cfg = yaml.safe_load(f)
77
+ dic = cfg.get("PostProcess", {}).get("character_dict", [])
78
+ if not dic:
79
+ raise ValueError(f"No PostProcess.character_dict found in {source}")
80
+ return dic
81
+ elif ext == ".txt":
82
+ with open(source, "r", encoding="utf-8") as f:
83
+ return [line.strip("\n\r") for line in f.readlines()]
84
+ else:
85
+ raise ValueError(f"Unsupported char_dict source: {source}. Use .yml, .txt, or list.")
86
+
87
+
88
+ # ============================================================================
89
+ # 1. Detection Preprocessing
90
+ # ============================================================================
91
+
92
+ class _DetResizeForTest:
93
+ def __init__(self, limit_side_len=960, limit_type="max", max_side_limit=4000):
94
+ self.limit_side_len = limit_side_len
95
+ self.limit_type = limit_type
96
+ self.max_side_limit = max_side_limit
97
+
98
+ def _image_padding(self, im, value=0):
99
+ h, w, c = im.shape
100
+ im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value
101
+ im_pad[:h, :w, :] = im
102
+ return im_pad
103
+
104
+ def _resize_image_type0(self, img):
105
+ h, w, _ = img.shape
106
+ limit_side_len = self.limit_side_len
107
+ if self.limit_type == "max":
108
+ ratio = float(limit_side_len) / max(h, w) if max(h, w) > limit_side_len else 1.0
109
+ elif self.limit_type == "min":
110
+ ratio = float(limit_side_len) / min(h, w) if min(h, w) < limit_side_len else 1.0
111
+ elif self.limit_type == "resize_long":
112
+ ratio = float(limit_side_len) / max(h, w)
113
+ else:
114
+ raise ValueError(f"not support limit_type: {self.limit_type}")
115
+
116
+ resize_h, resize_w = int(h * ratio), int(w * ratio)
117
+ if max(resize_h, resize_w) > self.max_side_limit:
118
+ ratio = float(self.max_side_limit) / max(resize_h, resize_w)
119
+ resize_h, resize_w = int(resize_h * ratio), int(resize_w * ratio)
120
+ resize_h = max(int(round(resize_h / 32) * 32), 32)
121
+ resize_w = max(int(round(resize_w / 32) * 32), 32)
122
+ if int(resize_w) <= 0 or int(resize_h) <= 0:
123
+ return None, (None, None)
124
+ img = cv2.resize(img, (int(resize_w), int(resize_h)))
125
+ ratio_h, ratio_w = resize_h / float(h), resize_w / float(w)
126
+ return img, [ratio_h, ratio_w]
127
+
128
+ def __call__(self, img):
129
+ src_h, src_w = img.shape[:2]
130
+ if sum([src_h, src_w]) < 64:
131
+ img = self._image_padding(img)
132
+ img, [ratio_h, ratio_w] = self._resize_image_type0(img)
133
+ shape = np.array([src_h, src_w, ratio_h, ratio_w])
134
+ return img, shape
135
+
136
+
137
+ class _NormalizeImage:
138
+ def __init__(self, mean, std, scale=1.0 / 255.0, order="hwc"):
139
+ self.scale = np.float32(scale)
140
+ shape = (1, 1, 3) if order == "hwc" else (3, 1, 1)
141
+ self.mean = np.array(mean, dtype=np.float32).reshape(shape)
142
+ self.std = np.array(std, dtype=np.float32).reshape(shape)
143
+
144
+ def __call__(self, img):
145
+ return (img.astype("float32") * self.scale - self.mean) / self.std
146
+
147
+
148
+ class _ToCHWImage:
149
+ def __call__(self, img):
150
+ return img.transpose((2, 0, 1))
151
+
152
+
153
+ # ============================================================================
154
+ # 2. Recognition Preprocessing
155
+ # ============================================================================
156
+
157
+ def _resize_norm_img(img, image_shape, max_wh_ratio=None):
158
+ imgC, imgH, imgW = image_shape
159
+ if max_wh_ratio is None:
160
+ max_wh_ratio = imgW * 1.0 / imgH
161
+ h, w = img.shape[:2]
162
+ max_wh_ratio = max(max_wh_ratio, w / h)
163
+ target_w = int(imgH * max_wh_ratio)
164
+ h, w = img.shape[:2]
165
+ ratio = w / h
166
+ resized_w = target_w if math.ceil(imgH * ratio) > target_w else int(math.ceil(imgH * ratio))
167
+ resized_image = cv2.resize(img, (resized_w, imgH)).astype("float32")
168
+ resized_image = resized_image.transpose((2, 0, 1)) / 255.0
169
+ resized_image -= 0.5
170
+ resized_image /= 0.5
171
+ padding_im = np.zeros((imgC, imgH, target_w), dtype=np.float32)
172
+ padding_im[:, :, 0:resized_w] = resized_image
173
+ return padding_im
174
+
175
+
176
+ # ============================================================================
177
+ # 3. Detection Postprocessing
178
+ # ============================================================================
179
+
180
+ class _DBPostProcess:
181
+ def __init__(self, thresh=0.3, box_thresh=0.7, max_candidates=1000,
182
+ unclip_ratio=2.0, use_dilation=False, score_mode="fast", box_type="quad"):
183
+ self.thresh = thresh
184
+ self.box_thresh = box_thresh
185
+ self.max_candidates = max_candidates
186
+ self.unclip_ratio = unclip_ratio
187
+ self.min_size = 3
188
+ self.score_mode = score_mode
189
+ self.box_type = box_type
190
+ self.dilation_kernel = None if not use_dilation else np.array([[1, 1], [1, 1]])
191
+
192
+ def _unclip(self, box, unclip_ratio):
193
+ poly = Polygon(box)
194
+ distance = poly.area * unclip_ratio / poly.length
195
+ offset = pyclipper.PyclipperOffset()
196
+ offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
197
+ return offset.Execute(distance)
198
+
199
+ def _get_mini_boxes(self, contour):
200
+ bb = cv2.minAreaRect(contour)
201
+ points = sorted(list(cv2.boxPoints(bb)), key=lambda x: x[0])
202
+ i1, i4 = (0, 1) if points[1][1] > points[0][1] else (1, 0)
203
+ i2, i3 = (2, 3) if points[3][1] > points[2][1] else (3, 2)
204
+ return [points[i1], points[i2], points[i3], points[i4]], min(bb[1])
205
+
206
+ def _box_score_fast(self, bitmap, _box):
207
+ h, w = bitmap.shape[:2]
208
+ box = _box.copy()
209
+ xmin = np.clip(np.floor(box[:, 0].min()).astype("int32"), 0, w - 1)
210
+ xmax = np.clip(np.ceil(box[:, 0].max()).astype("int32"), 0, w - 1)
211
+ ymin = np.clip(np.floor(box[:, 1].min()).astype("int32"), 0, h - 1)
212
+ ymax = np.clip(np.ceil(box[:, 1].max()).astype("int32"), 0, h - 1)
213
+ mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
214
+ box[:, 0] -= xmin
215
+ box[:, 1] -= ymin
216
+ cv2.fillPoly(mask, box.reshape(1, -1, 2).astype("int32"), 1)
217
+ return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
218
+
219
+ def _boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
220
+ height, width = _bitmap.shape
221
+ outs = cv2.findContours(
222
+ (_bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
223
+ )
224
+ contours = outs[0] if len(outs) == 2 else outs[1]
225
+ num = min(len(contours), self.max_candidates)
226
+ boxes, scores = [], []
227
+ for i in range(num):
228
+ points, sside = self._get_mini_boxes(contours[i])
229
+ if sside < self.min_size:
230
+ continue
231
+ points = np.array(points)
232
+ score = self._box_score_fast(pred, points.reshape(-1, 2))
233
+ if self.box_thresh > score:
234
+ continue
235
+ box = self._unclip(points, self.unclip_ratio)
236
+ if len(box) > 1:
237
+ continue
238
+ box = np.array(box).reshape(-1, 1, 2)
239
+ box, sside = self._get_mini_boxes(box)
240
+ if sside < self.min_size + 2:
241
+ continue
242
+ box = np.array(box)
243
+ box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
244
+ box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
245
+ boxes.append(box.astype("int32"))
246
+ scores.append(score)
247
+ return np.array(boxes, dtype="int32"), scores
248
+
249
+ def __call__(self, pred, shape_list):
250
+ pred = pred[:, 0, :, :]
251
+ segmentation = pred > self.thresh
252
+ boxes_batch = []
253
+ for bi in range(pred.shape[0]):
254
+ src_h, src_w, ratio_h, ratio_w = shape_list[bi]
255
+ mask = cv2.dilate(np.array(segmentation[bi]).astype(np.uint8),
256
+ self.dilation_kernel) if self.dilation_kernel is not None else segmentation[bi]
257
+ boxes, _ = self._boxes_from_bitmap(pred[bi], mask, src_w, src_h)
258
+ boxes_batch.append(boxes)
259
+ return boxes_batch
260
+
261
+
262
+ # ============================================================================
263
+ # 4. Recognition Postprocessing
264
+ # ============================================================================
265
+
266
+ class _CTCLabelDecode:
267
+ def __init__(self, character_list: List[str], use_space_char=True):
268
+ self.character_str = list(character_list)
269
+ if use_space_char:
270
+ self.character_str.append(" ")
271
+ dict_character = ["blank"] + self.character_str
272
+ self.character = dict_character
273
+
274
+ def decode(self, text_index, text_prob=None, is_remove_duplicate=True):
275
+ result_list = []
276
+ for bi in range(len(text_index)):
277
+ sel = np.ones(len(text_index[bi]), dtype=bool)
278
+ if is_remove_duplicate:
279
+ sel[1:] = text_index[bi][1:] != text_index[bi][:-1]
280
+ sel &= text_index[bi] != 0
281
+ chars = [self.character[int(t)] for t in text_index[bi][sel]]
282
+ conf = text_prob[bi][sel] if text_prob is not None else [1] * len(sel)
283
+ if len(conf) == 0:
284
+ conf = [0]
285
+ result_list.append(("".join(chars), float(np.mean(conf))))
286
+ return result_list
287
+
288
+ def __call__(self, preds):
289
+ return self.decode(preds.argmax(axis=2), preds.max(axis=2), is_remove_duplicate=True)
290
+
291
+
292
+ # ============================================================================
293
+ # 5. Image Utilities
294
+ # ============================================================================
295
+
296
+ def _get_rotate_crop_image(img: np.ndarray, points: np.ndarray) -> np.ndarray:
297
+ assert len(points) == 4
298
+ cw = int(max(np.linalg.norm(points[0] - points[1]), np.linalg.norm(points[2] - points[3])))
299
+ ch = int(max(np.linalg.norm(points[0] - points[3]), np.linalg.norm(points[1] - points[2])))
300
+ pts_std = np.float32([[0, 0], [cw, 0], [cw, ch], [0, ch]])
301
+ M = cv2.getPerspectiveTransform(points.astype(np.float32), pts_std)
302
+ dst = cv2.warpPerspective(img, M, (cw, ch), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC)
303
+ if dst.shape[0] * 1.0 / dst.shape[1] >= 1.5:
304
+ dst = np.rot90(dst)
305
+ return dst
306
+
307
+
308
+ def _sorted_boxes(dt_boxes):
309
+ if len(dt_boxes) == 0:
310
+ return dt_boxes
311
+ boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
312
+ lst = list(boxes)
313
+ for i in range(len(lst) - 1):
314
+ for j in range(i, -1, -1):
315
+ if abs(lst[j + 1][0][1] - lst[j][0][1]) < 10 and lst[j + 1][0][0] < lst[j][0][0]:
316
+ lst[j], lst[j + 1] = lst[j + 1], lst[j]
317
+ else:
318
+ break
319
+ return lst
320
+
321
+
322
+ def draw_ocr_result(
323
+ img: np.ndarray,
324
+ results: List[dict],
325
+ font_path: str = "./fonts/simfang.ttf",
326
+ ) -> np.ndarray:
327
+ """Draw detection boxes (semi-transparent) on original image, with text list on the right side."""
328
+ h, w = img.shape[:2]
329
+
330
+ # --- left: original image with semi-transparent colored boxes ---
331
+ pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
332
+ overlay = Image.new("RGBA", pil_img.size, (0, 0, 0, 0))
333
+ draw_overlay = ImageDraw.Draw(overlay)
334
+
335
+ random.seed(0)
336
+ for res in results:
337
+ box = res["box"]
338
+ color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 90)
339
+ draw_overlay.polygon([tuple(p) for p in box], fill=color)
340
+
341
+ left_img = Image.alpha_composite(pil_img.convert("RGBA"), overlay).convert("RGB")
342
+
343
+ # --- right: white canvas with text list ---
344
+ right_w = int(w * 0.9)
345
+ right = Image.new("RGB", (right_w, h), (255, 255, 255))
346
+ draw_right = ImageDraw.Draw(right)
347
+
348
+ try:
349
+ font = ImageFont.truetype(font_path, 14)
350
+ except (OSError, IOError):
351
+ font = ImageFont.load_default()
352
+
353
+ y = 5
354
+ gap = 18
355
+ for i, res in enumerate(results):
356
+ text = f"{i+1}. {res['text']} ({res['confidence']:.3f})"
357
+ # Color block matching the box
358
+ random.seed(i)
359
+ blk_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
360
+ draw_right.rectangle([5, y + 3, 15, y + 14], fill=blk_color, outline=(0, 0, 0))
361
+ draw_right.text((20, y), text, fill=(0, 0, 0), font=font)
362
+ y += gap
363
+
364
+ # --- concat left + right ---
365
+ result_img = Image.new("RGB", (w + right_w, h))
366
+ result_img.paste(left_img, (0, 0))
367
+ result_img.paste(right, (w, 0))
368
+
369
+ return cv2.cvtColor(np.array(result_img), cv2.COLOR_RGB2BGR)
370
+
371
+
372
+ # ============================================================================
373
+ # 6. Main Inference Engine
374
+ # ============================================================================
375
+
376
+ class PPOCRv6Onnx:
377
+
378
+ def __init__(
379
+ self,
380
+ det_onnx: str,
381
+ rec_onnx: str,
382
+ char_dict: Union[str, List[str]],
383
+ # Detection params
384
+ det_limit_side_len: int = 960,
385
+ det_db_thresh: float = 0.2,
386
+ det_db_box_thresh: float = 0.4,
387
+ det_db_unclip_ratio: float = 1.4,
388
+ det_max_candidates: int = 3000,
389
+ # Recognition params
390
+ rec_image_shape: Tuple[int, int, int] = (3, 48, 320),
391
+ rec_batch_num: int = 6,
392
+ # Classifier params
393
+ use_angle_cls: bool = False,
394
+ cls_onnx: Optional[str] = None,
395
+ cls_image_shape: Tuple[int, int, int] = (3, 48, 192),
396
+ cls_batch_num: int = 6,
397
+ cls_thresh: float = 0.9,
398
+ cls_label_list: Optional[List[str]] = None,
399
+ # Common
400
+ drop_score: float = 0.5,
401
+ use_gpu: bool = False,
402
+ onnx_providers: Optional[List[str]] = None,
403
+ resize_mode: str = "letterbox",
404
+ ):
405
+ assert resize_mode in ("letterbox", "stretch"), f"invalid resize_mode: {resize_mode}"
406
+ if cls_label_list is None:
407
+ cls_label_list = ["0", "180"]
408
+
409
+ self.rec_image_shape = rec_image_shape
410
+ self.rec_batch_num = rec_batch_num
411
+ self.drop_score = drop_score
412
+ self.use_angle_cls = use_angle_cls
413
+ self.cls_thresh = cls_thresh
414
+ self.cls_label_list = cls_label_list
415
+ self.cls_batch_num = cls_batch_num
416
+ self.cls_image_shape = cls_image_shape
417
+ self._resize_mode = resize_mode
418
+
419
+ # ---- ONNX sessions ----
420
+ if onnx_providers is None:
421
+ onnx_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if use_gpu else ["CPUExecutionProvider"]
422
+ sess_options = ort.SessionOptions()
423
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
424
+
425
+ self.det_session = ort.InferenceSession(det_onnx, sess_options=sess_options, providers=onnx_providers)
426
+ self.det_input_name = self.det_session.get_inputs()[0].name
427
+
428
+ self.rec_session = ort.InferenceSession(rec_onnx, sess_options=sess_options, providers=onnx_providers)
429
+ self.rec_input_name = self.rec_session.get_inputs()[0].name
430
+
431
+ # Classifier session
432
+ if use_angle_cls:
433
+ if cls_onnx is None:
434
+ raise ValueError("cls_onnx is required when use_angle_cls=True")
435
+ self.cls_session = ort.InferenceSession(cls_onnx, sess_options=sess_options, providers=onnx_providers)
436
+ self.cls_input_name = self.cls_session.get_inputs()[0].name
437
+ # Detect fixed dims for cls
438
+ cls_h, cls_w = _detect_fixed_dims(self.cls_session, cls_onnx)
439
+ self._cls_fixed_h = cls_h if cls_h > 0 else 0
440
+ self._cls_fixed_w = cls_w if cls_w > 0 else 0
441
+ else:
442
+ self.cls_session = None
443
+
444
+ # Detect fixed dims for det/rec
445
+ self._det_fixed_h, self._det_fixed_w = _detect_fixed_dims(self.det_session, det_onnx)
446
+ det_shape = self.det_session.get_inputs()[0].shape
447
+ print(f"[PPOCRv6] det shape={det_shape}, fixed_h={self._det_fixed_h}, fixed_w={self._det_fixed_w}, "
448
+ f"cls={use_angle_cls}, resize_mode={resize_mode}")
449
+
450
+ # Fixed rec width from ONNX
451
+ rec_inp = self.rec_session.get_inputs()[0]
452
+ rec_fw = _get_dim_value(rec_inp.shape[3])
453
+ self._rec_fixed_w = rec_fw if rec_fw > 0 else 0
454
+
455
+ # ---- Detection pre/post ----
456
+ self._det_resize = _DetResizeForTest(limit_side_len=det_limit_side_len, limit_type="max")
457
+ self._det_normalize = _NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
458
+ self._det_to_chw = _ToCHWImage()
459
+ self._det_post = _DBPostProcess(thresh=det_db_thresh, box_thresh=det_db_box_thresh,
460
+ unclip_ratio=det_db_unclip_ratio, max_candidates=det_max_candidates, box_type="quad")
461
+
462
+ # ---- Recognition post ----
463
+ self._rec_post = _CTCLabelDecode(_load_char_dict(char_dict), use_space_char=True)
464
+
465
+ # ---- Detection with static shape support ----
466
+
467
+ def _preprocess_det(self, img: np.ndarray):
468
+ src_h, src_w = img.shape[:2]
469
+ fh, fw = self._det_fixed_h, self._det_fixed_w
470
+
471
+ if fh > 0 and fw > 0:
472
+ if self._resize_mode == "stretch":
473
+ # Direct resize (official behavior)
474
+ img_r = cv2.resize(img, (fw, fh))
475
+ shape = np.array([src_h, src_w, float(fh) / src_h, float(fw) / src_w])
476
+ else:
477
+ # Letterbox: ratio-preserving + pad
478
+ ratio = min(fh / src_h, fw / src_w)
479
+ new_h = min(max(int(round(src_h * ratio / 32) * 32), 32), fh)
480
+ new_w = min(max(int(round(src_w * ratio / 32) * 32), 32), fw)
481
+ img_r = cv2.resize(img, (new_w, new_h))
482
+ pad_h, pad_w = max(0, fh - new_h), max(0, fw - new_w)
483
+ if pad_h or pad_w:
484
+ img_r = cv2.copyMakeBorder(img_r, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=(0, 0, 0))
485
+ shape = np.array([src_h * fh / new_h, src_w * fw / new_w,
486
+ float(new_h) / src_h, float(new_w) / src_w])
487
+ else:
488
+ img_r, shape = self._det_resize(img)
489
+
490
+ img_n = self._det_normalize(img_r)
491
+ img_c = self._det_to_chw(img_n)
492
+ return np.expand_dims(img_c.astype(np.float32), axis=0), shape
493
+
494
+ def _postprocess_det(self, output, shape):
495
+ return self._det_post(output, np.expand_dims(shape, axis=0))[0]
496
+
497
+ def detect(self, img):
498
+ tensor, shape = self._preprocess_det(img)
499
+ out = self.det_session.run(None, {self.det_input_name: tensor})
500
+ return self._postprocess_det(out[0], shape)
501
+
502
+ # ---- Classifier ----
503
+
504
+ def _preprocess_cls(self, img_list):
505
+ num = len(img_list)
506
+ width_list = [im.shape[1] / float(im.shape[0]) for im in img_list]
507
+ indices = np.argsort(np.array(width_list))
508
+ batches, idx_maps = [], []
509
+ for beg in range(0, num, self.cls_batch_num):
510
+ end = min(num, beg + self.cls_batch_num)
511
+ imgC, imgH, imgW = self.cls_image_shape
512
+
513
+ # Override with ONNX fixed dimensions
514
+ if self._cls_fixed_h > 0:
515
+ imgH = self._cls_fixed_h
516
+ if self._cls_fixed_w > 0:
517
+ imgW = self._cls_fixed_w
518
+
519
+ max_wh_ratio = imgW / imgH
520
+ for ino in range(beg, end):
521
+ h, w = img_list[indices[ino]].shape[:2]
522
+ max_wh_ratio = max(max_wh_ratio, w / h)
523
+ if self._cls_fixed_w > 0:
524
+ max_wh_ratio = self._cls_fixed_w / imgH
525
+
526
+ shape = (imgC, imgH, imgW)
527
+ norm_list, idx_list = [], []
528
+ for ino in range(beg, end):
529
+ orig_idx = indices[ino]
530
+ norm = _resize_norm_img(img_list[orig_idx], shape, max_wh_ratio=max_wh_ratio)
531
+ norm_list.append(np.expand_dims(norm, axis=0))
532
+ idx_list.append(orig_idx)
533
+ if norm_list:
534
+ batches.append(np.concatenate(norm_list, axis=0).astype(np.float32))
535
+ idx_maps.append(idx_list)
536
+ return batches, idx_maps
537
+
538
+ def _postprocess_cls(self, batch_outputs, idx_maps, total_num, img_list):
539
+ results = [("0", 1.0)] * total_num
540
+ for preds_batch, idx_list in zip(batch_outputs, idx_maps):
541
+ pred_ids = preds_batch.argmax(axis=1)
542
+ for i, orig_idx in enumerate(idx_list):
543
+ label = self.cls_label_list[int(pred_ids[i])]
544
+ score = float(preds_batch[i, int(pred_ids[i])])
545
+ results[orig_idx] = (label, score)
546
+ if "180" in str(label) and score > self.cls_thresh:
547
+ img_list[orig_idx] = cv2.rotate(img_list[orig_idx], cv2.ROTATE_180)
548
+ return results
549
+
550
+ def classify(self, img_list):
551
+ if not img_list or not self.use_angle_cls:
552
+ return img_list, [], 0
553
+ img_list = [im.copy() for im in img_list]
554
+ batches, idx_maps = self._preprocess_cls(img_list)
555
+ outputs = []
556
+ for batch in batches:
557
+ out = self.cls_session.run(None, {self.cls_input_name: batch})
558
+ outputs.append(out[0])
559
+ cls_res = self._postprocess_cls(outputs, idx_maps, len(img_list), img_list)
560
+ return img_list, cls_res, 0
561
+
562
+ # ---- Recognition ----
563
+
564
+ def _preprocess_rec(self, img_crop_list):
565
+ num = len(img_crop_list)
566
+ width_list = [im.shape[1] / float(im.shape[0]) for im in img_crop_list]
567
+ indices = np.argsort(np.array(width_list))
568
+ batches, idx_maps = [], []
569
+ for beg in range(0, num, self.rec_batch_num):
570
+ end = min(num, beg + self.rec_batch_num)
571
+ imgC, imgH, imgW = self.rec_image_shape
572
+ max_wh_ratio = imgW / imgH
573
+ for ino in range(beg, end):
574
+ h, w = img_crop_list[indices[ino]].shape[:2]
575
+ max_wh_ratio = max(max_wh_ratio, w / h)
576
+ # Use fixed width if set by ONNX
577
+ if self._rec_fixed_w > 0:
578
+ max_wh_ratio = self._rec_fixed_w / imgH
579
+ norm_list, idx_list = [], []
580
+ for ino in range(beg, end):
581
+ orig_idx = indices[ino]
582
+ norm = _resize_norm_img(img_crop_list[orig_idx], self.rec_image_shape, max_wh_ratio=max_wh_ratio)
583
+ norm_list.append(np.expand_dims(norm, axis=0))
584
+ idx_list.append(orig_idx)
585
+ if norm_list:
586
+ batches.append(np.concatenate(norm_list, axis=0).astype(np.float32))
587
+ idx_maps.append(idx_list)
588
+ return batches, idx_maps
589
+
590
+ def _postprocess_rec(self, batch_outputs, idx_maps, total_num):
591
+ results = [("", 0.0)] * total_num
592
+ # Decode each batch separately (different T per batch)
593
+ for preds_batch, idx_list in zip(batch_outputs, idx_maps):
594
+ texts = self._rec_post(preds_batch)
595
+ for i, orig_idx in enumerate(idx_list):
596
+ results[orig_idx] = texts[i]
597
+ return results
598
+
599
+ def recognize(self, img_crop_list):
600
+ if not img_crop_list:
601
+ return []
602
+ batches, idx_maps = self._preprocess_rec(img_crop_list)
603
+ outputs = []
604
+ for batch in batches:
605
+ out = self.rec_session.run(None, {self.rec_input_name: batch})
606
+ outputs.append(out[0])
607
+ return self._postprocess_rec(outputs, idx_maps, len(img_crop_list))
608
+
609
+ # ---- Full pipeline ----
610
+
611
+ def predict_image(self, image_path, visualize=False):
612
+ img = cv2.imread(image_path)
613
+ if img is None:
614
+ raise FileNotFoundError(f"Cannot read image: {image_path}")
615
+ return self(img, visualize=visualize)
616
+
617
+ def __call__(self, img: np.ndarray, visualize=False, use_cls=None):
618
+ ori_im = img.copy()
619
+
620
+ boxes = self.detect(img)
621
+ if len(boxes) == 0:
622
+ return [] if not visualize else ori_im
623
+
624
+ boxes = _sorted_boxes(boxes)
625
+
626
+ img_crop_list = []
627
+ for box in boxes:
628
+ crop = _get_rotate_crop_image(ori_im, np.array(box, dtype=np.float32))
629
+ img_crop_list.append(crop)
630
+
631
+ # Direction classifier
632
+ do_cls = self.use_angle_cls if use_cls is None else use_cls
633
+ if do_cls and self.cls_session is not None:
634
+ img_crop_list, cls_res, _ = self.classify(img_crop_list)
635
+
636
+ rec_res = self.recognize(img_crop_list)
637
+
638
+ results = []
639
+ for box, (text, conf) in zip(boxes, rec_res):
640
+ if conf >= self.drop_score:
641
+ results.append({"text": text, "confidence": round(conf, 4), "box": box.tolist()})
642
+
643
+ return draw_ocr_result(ori_im, results) if visualize else results
644
+
645
+
646
+ # ============================================================================
647
+ # 7. CLI
648
+ # ============================================================================
649
+
650
+ def main():
651
+ parser = argparse.ArgumentParser(description="PP-OCRv6 ONNX Inference (standalone, no Paddle dependency)")
652
+ parser.add_argument("--det_onnx", type=str, default="onnx/det_inference_static_sim.onnx")
653
+ parser.add_argument("--rec_onnx", type=str, default="onnx/rec_inference_static_sim.onnx")
654
+ parser.add_argument("--char_dict", type=str, default="onnx/rec_inference.yml")
655
+ parser.add_argument("--image", required=True, help="Input image path")
656
+ parser.add_argument("--use_gpu", action="store_true")
657
+ parser.add_argument("--drop_score", type=float, default=0.5)
658
+ parser.add_argument("--det_limit_side_len", type=int, default=960)
659
+ parser.add_argument("--det_db_thresh", type=float, default=0.2)
660
+ parser.add_argument("--det_db_box_thresh", type=float, default=0.45)
661
+ parser.add_argument("--det_db_unclip_ratio", type=float, default=1.4)
662
+ parser.add_argument("--rec_batch_num", type=int, default=1)
663
+ parser.add_argument("--resize_mode", type=str, default="letterbox", choices=["letterbox", "stretch"])
664
+ # Classifier
665
+ parser.add_argument("--use_angle_cls", action="store_true", help="Enable direction classifier")
666
+ parser.add_argument("--cls_onnx", type=str, default="onnx/cls_mobile_sim_static.onnx", help="Classifer ONNX model path")
667
+ parser.add_argument("--cls_thresh", type=float, default=0.9, help="Angie classifier confidence threshold")
668
+ parser.add_argument("--cls_batch_num", type=int, default=1)
669
+ # Output
670
+ parser.add_argument("--visualize", action="store_true")
671
+ parser.add_argument("--output", type=str, default=None)
672
+ parser.add_argument("--json", type=str, default=None)
673
+
674
+ args = parser.parse_args()
675
+
676
+ char_dict_src = args.char_dict
677
+ if not os.path.exists(char_dict_src) and ("," in char_dict_src or char_dict_src.startswith("[")):
678
+ char_dict = [c.strip() for c in char_dict_src.strip("[]").split(",") if c.strip()]
679
+ else:
680
+ char_dict = char_dict_src
681
+
682
+ ocr = PPOCRv6Onnx(
683
+ det_onnx=args.det_onnx,
684
+ rec_onnx=args.rec_onnx,
685
+ char_dict=char_dict,
686
+ det_limit_side_len=args.det_limit_side_len,
687
+ det_db_thresh=args.det_db_thresh,
688
+ det_db_box_thresh=args.det_db_box_thresh,
689
+ det_db_unclip_ratio=args.det_db_unclip_ratio,
690
+ rec_batch_num=args.rec_batch_num,
691
+ use_angle_cls=args.use_angle_cls,
692
+ cls_onnx=args.cls_onnx,
693
+ cls_thresh=args.cls_thresh,
694
+ cls_batch_num=args.cls_batch_num,
695
+ drop_score=args.drop_score,
696
+ use_gpu=args.use_gpu,
697
+ resize_mode=args.resize_mode,
698
+ )
699
+
700
+ do_viz = args.visualize or args.output is not None
701
+ img = cv2.imread(args.image)
702
+ if img is None:
703
+ raise FileNotFoundError(f"Cannot read image: {args.image}")
704
+ results = ocr(img)
705
+
706
+ if do_viz:
707
+ vis = ocr(img, visualize=True)
708
+ out_path = args.output or "res-onnx.jpg"
709
+ cv2.imwrite(out_path, vis)
710
+ print(f"Annotated image saved to: {out_path}")
711
+
712
+ # if args.json:
713
+ # with open(args.json, "w", encoding="utf-8") as f:
714
+ # json.dump(results, f, ensure_ascii=False, indent=2)
715
+ # print(f"Results saved to: {args.json}")
716
+ # else:
717
+ # print(json.dumps(results, ensure_ascii=False, indent=2))
718
+
719
+
720
+ if __name__ == "__main__":
721
+ main()
rec.json ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "input": "./onnx/rec_inference_static_sim.onnx",
3
+ "output_dir": "./rec-ax650-npu1",
4
+ "output_name": "rec_npu1.axmodel",
5
+ "model_type": "ONNX",
6
+ "target_hardware": "AX650",
7
+ "npu_mode": "NPU1",
8
+
9
+ // "onnx_opt": {
10
+ // "disable_onnx_optimization": false,
11
+ // "enable_onnxsim": false,
12
+ // "model_check": true,
13
+ // "disable_transformation_check": false
14
+ // },
15
+ "quant": {
16
+ "input_configs": [
17
+ {
18
+ "tensor_name": "DEFAULT",
19
+ "calibration_dataset": "./dataset/rec_quant.tar",
20
+ "calibration_format": "Image",
21
+ "calibration_size": -1,
22
+ "calibration_mean": [127.5, 127.5, 127.5],
23
+ "calibration_std": [127.5, 127.5, 127.5]
24
+ }
25
+ ],
26
+
27
+ "calibration_method": "MinMax",
28
+ "precision_analysis": true,
29
+ "precision_analysis_method": "EndToEnd",
30
+ "precision_analysis_mode": "Reference",
31
+ "conv_bias_data_type": "FP32",
32
+ "layer_configs": [
33
+ {
34
+ "start_tensor_names": ["DEFAULT"],
35
+ "end_tensor_names": ["DEFAULT"],
36
+ "data_type": "U16"
37
+ }
38
+ ]
39
+ },
40
+ "input_processors": [
41
+ {
42
+ "tensor_name": "DEFAULT",
43
+ "tensor_format": "BGR",
44
+ "tensor_layout": "NCHW",
45
+ "src_format": "BGR",
46
+ "src_dtype": "FP32",
47
+ "src_layout": "NCHW",
48
+ "csc_mode": "NoCSC"
49
+ }
50
+ ],
51
+ "compiler": {
52
+ "check": 2
53
+ }
54
+ }
res-ax.jpg ADDED

Git LFS Details

  • SHA256: 5ba25dc987e45e63b7e02293ca90e7f90c2f8ed170c6fd2905cec42de77ad505
  • Pointer size: 131 Bytes
  • Size of remote file: 140 kB
res-onnx.jpg ADDED

Git LFS Details

  • SHA256: f02b26ff40163f9c3917d0210cc5b4f996f54b6937c9427776e59b274eebb280
  • Pointer size: 131 Bytes
  • Size of remote file: 140 kB
run_det_ax.py ADDED
@@ -0,0 +1,894 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ PP-OCRv6 Detection ONNX Inference & Evaluation (standalone, zero Paddle dependency)
4
+
5
+ Dependencies:
6
+ numpy, opencv-python, onnxruntime, pyyaml, shapely, pyclipper
7
+
8
+ Data format (same as PaddleOCR official):
9
+ Label file: image_path<TAB>json_label
10
+ json_label: [{"transcription": "text_or_###", "points": [[x,y]*4]}, ...]
11
+
12
+ Usage:
13
+ # Single image inference
14
+ python ppocrv6_det_onnx.py --det_onnx det.onnx --image test.jpg --visualize
15
+
16
+ # Batch evaluation
17
+ python ppocrv6_det_onnx.py --det_onnx det.onnx \\
18
+ --label_file val.txt --dataset_root ./ocr_det_dataset_examples \\
19
+ --visualize --output_json result.json --verbose
20
+ """
21
+
22
+ import argparse
23
+ import json
24
+ import math
25
+ import os
26
+ from collections import namedtuple
27
+ from typing import List, Optional, Tuple
28
+
29
+ import cv2
30
+ import numpy as np
31
+ import axengine as ort
32
+ import yaml
33
+ from shapely.geometry import Polygon
34
+ import pyclipper
35
+
36
+
37
+ def _get_dim_value(dim):
38
+ """Extract integer value from an ONNX Runtime dimension, returning 0 for dynamic dims."""
39
+ if dim is None:
40
+ return 0
41
+ if isinstance(dim, str):
42
+ return 0
43
+ if hasattr(dim, 'dim_value'):
44
+ return int(dim.dim_value) if dim.dim_value else 0
45
+ if hasattr(dim, 'dim_param'):
46
+ # named dim like "batch_size" → dynamic
47
+ return 0
48
+ try:
49
+ v = int(dim)
50
+ return v
51
+ except (TypeError, ValueError):
52
+ return 0
53
+
54
+
55
+ # ============================================================================
56
+ # 1. Detection Preprocessing
57
+ # ============================================================================
58
+
59
+ class _DetResizeForTest:
60
+ def __init__(self, limit_side_len=960, limit_type="max", max_side_limit=4000,
61
+ image_shape=None, keep_ratio=False):
62
+ self.max_side_limit = max_side_limit
63
+ # resize_type=0: limit_side_len (dynamic input)
64
+ # resize_type=1: image_shape (fixed ONNX input)
65
+ if image_shape is not None:
66
+ self.resize_type = 1
67
+ self.image_shape = image_shape
68
+ self.keep_ratio = keep_ratio
69
+ else:
70
+ self.resize_type = 0
71
+ self.limit_side_len = limit_side_len
72
+ self.limit_type = limit_type
73
+
74
+ def _image_padding(self, im, value=0):
75
+ h, w, c = im.shape
76
+ im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value
77
+ im_pad[:h, :w, :] = im
78
+ return im_pad
79
+
80
+ def _resize_image_type0(self, img):
81
+ h, w, _ = img.shape
82
+ limit_side_len = self.limit_side_len
83
+
84
+ if self.limit_type == "max":
85
+ if max(h, w) > limit_side_len:
86
+ ratio = float(limit_side_len) / max(h, w)
87
+ else:
88
+ ratio = 1.0
89
+ elif self.limit_type == "min":
90
+ if min(h, w) < limit_side_len:
91
+ ratio = float(limit_side_len) / min(h, w)
92
+ else:
93
+ ratio = 1.0
94
+ elif self.limit_type == "resize_long":
95
+ ratio = float(limit_side_len) / max(h, w)
96
+ else:
97
+ raise ValueError(f"not support limit_type: {self.limit_type}")
98
+
99
+ resize_h = int(h * ratio)
100
+ resize_w = int(w * ratio)
101
+ if max(resize_h, resize_w) > self.max_side_limit:
102
+ ratio = float(self.max_side_limit) / max(resize_h, resize_w)
103
+ resize_h, resize_w = int(resize_h * ratio), int(resize_w * ratio)
104
+
105
+ resize_h = max(int(round(resize_h / 32) * 32), 32)
106
+ resize_w = max(int(round(resize_w / 32) * 32), 32)
107
+
108
+ if int(resize_w) <= 0 or int(resize_h) <= 0:
109
+ return None, (None, None)
110
+ img = cv2.resize(img, (int(resize_w), int(resize_h)))
111
+ ratio_h = resize_h / float(h)
112
+ ratio_w = resize_w / float(w)
113
+ return img, [ratio_h, ratio_w]
114
+
115
+ def _resize_image_type1(self, img):
116
+ """Direct resize to fixed [H, W]. Used when ONNX has fixed input dimensions."""
117
+ resize_h, resize_w = self.image_shape
118
+ ori_h, ori_w = img.shape[:2]
119
+ if self.keep_ratio:
120
+ resize_w = ori_w * resize_h / ori_h
121
+ N = math.ceil(resize_w / 32)
122
+ resize_w = N * 32
123
+ ratio_h = float(resize_h) / ori_h
124
+ ratio_w = float(resize_w) / ori_w
125
+ img = cv2.resize(img, (int(resize_w), int(resize_h)))
126
+ return img, [ratio_h, ratio_w]
127
+
128
+ def __call__(self, img):
129
+ src_h, src_w = img.shape[:2]
130
+ if sum([src_h, src_w]) < 64:
131
+ img = self._image_padding(img)
132
+ if self.resize_type == 1:
133
+ img, [ratio_h, ratio_w] = self._resize_image_type1(img)
134
+ else:
135
+ img, [ratio_h, ratio_w] = self._resize_image_type0(img)
136
+ shape = np.array([src_h, src_w, ratio_h, ratio_w])
137
+ return img, shape
138
+
139
+
140
+ class _NormalizeImage:
141
+ def __init__(self, mean, std, scale=1.0 / 255.0, order="hwc"):
142
+ self.scale = np.float32(scale)
143
+ shape = (1, 1, 3) if order == "hwc" else (3, 1, 1)
144
+ self.mean = np.array(mean, dtype=np.float32).reshape(shape)
145
+ self.std = np.array(std, dtype=np.float32).reshape(shape)
146
+
147
+ def __call__(self, img):
148
+ return (img.astype("float32") * self.scale - self.mean) / self.std
149
+
150
+
151
+ class _ToCHWImage:
152
+ def __call__(self, img):
153
+ return img.transpose((2, 0, 1))
154
+
155
+
156
+ # ============================================================================
157
+ # 2. Detection Postprocessing (DB)
158
+ # ============================================================================
159
+
160
+ class _DBPostProcess:
161
+ def __init__(
162
+ self,
163
+ thresh=0.3,
164
+ box_thresh=0.7,
165
+ max_candidates=1000,
166
+ unclip_ratio=2.0,
167
+ use_dilation=False,
168
+ score_mode="fast",
169
+ ):
170
+ self.thresh = thresh
171
+ self.box_thresh = box_thresh
172
+ self.max_candidates = max_candidates
173
+ self.unclip_ratio = unclip_ratio
174
+ self.min_size = 3
175
+ self.score_mode = score_mode
176
+ assert score_mode in ("slow", "fast")
177
+ self.dilation_kernel = None if not use_dilation else np.array([[1, 1], [1, 1]])
178
+
179
+ def _unclip(self, box, unclip_ratio):
180
+ poly = Polygon(box)
181
+ distance = poly.area * unclip_ratio / poly.length
182
+ offset = pyclipper.PyclipperOffset()
183
+ offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
184
+ expanded = offset.Execute(distance)
185
+ return expanded
186
+
187
+ def _get_mini_boxes(self, contour):
188
+ bounding_box = cv2.minAreaRect(contour)
189
+ points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
190
+ i1, i2, i3, i4 = 0, 1, 2, 3
191
+ if points[1][1] > points[0][1]:
192
+ i1, i4 = 0, 1
193
+ else:
194
+ i1, i4 = 1, 0
195
+ if points[3][1] > points[2][1]:
196
+ i2, i3 = 2, 3
197
+ else:
198
+ i2, i3 = 3, 2
199
+ box = [points[i1], points[i2], points[i3], points[i4]]
200
+ return box, min(bounding_box[1])
201
+
202
+ def _box_score_fast(self, bitmap, _box):
203
+ h, w = bitmap.shape[:2]
204
+ box = _box.copy()
205
+ xmin = np.clip(np.floor(box[:, 0].min()).astype("int32"), 0, w - 1)
206
+ xmax = np.clip(np.ceil(box[:, 0].max()).astype("int32"), 0, w - 1)
207
+ ymin = np.clip(np.floor(box[:, 1].min()).astype("int32"), 0, h - 1)
208
+ ymax = np.clip(np.ceil(box[:, 1].max()).astype("int32"), 0, h - 1)
209
+ mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
210
+ box[:, 0] = box[:, 0] - xmin
211
+ box[:, 1] = box[:, 1] - ymin
212
+ cv2.fillPoly(mask, box.reshape(1, -1, 2).astype("int32"), 1)
213
+ return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]
214
+
215
+ def _boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
216
+ bitmap = _bitmap
217
+ height, width = bitmap.shape
218
+ outs = cv2.findContours(
219
+ (bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
220
+ )
221
+ if len(outs) == 3:
222
+ _, contours, _ = outs
223
+ else:
224
+ contours, _ = outs
225
+ num_contours = min(len(contours), self.max_candidates)
226
+ boxes, scores = [], []
227
+ for index in range(num_contours):
228
+ contour = contours[index]
229
+ points, sside = self._get_mini_boxes(contour)
230
+ if sside < self.min_size:
231
+ continue
232
+ points = np.array(points)
233
+ if self.score_mode == "fast":
234
+ score = self._box_score_fast(pred, points.reshape(-1, 2))
235
+ else:
236
+ score = self._box_score_slow(pred, contour)
237
+ if self.box_thresh > score:
238
+ continue
239
+ box = self._unclip(points, self.unclip_ratio)
240
+ if len(box) > 1:
241
+ continue
242
+ box = np.array(box).reshape(-1, 1, 2)
243
+ box, sside = self._get_mini_boxes(box)
244
+ if sside < self.min_size + 2:
245
+ continue
246
+ box = np.array(box)
247
+ box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
248
+ box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
249
+ boxes.append(box.astype("int32"))
250
+ scores.append(score)
251
+ return np.array(boxes, dtype="int32"), scores
252
+
253
+ def _box_score_slow(self, bitmap, contour):
254
+ h, w = bitmap.shape[:2]
255
+ contour = contour.copy().reshape((-1, 2))
256
+ xmin = np.clip(np.min(contour[:, 0]), 0, w - 1)
257
+ xmax = np.clip(np.max(contour[:, 0]), 0, w - 1)
258
+ ymin = np.clip(np.min(contour[:, 1]), 0, h - 1)
259
+ ymax = np.clip(np.max(contour[:, 1]), 0, h - 1)
260
+ mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
261
+ contour[:, 0] = contour[:, 0] - xmin
262
+ contour[:, 1] = contour[:, 1] - ymin
263
+ cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype("int32"), 1)
264
+ return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]
265
+
266
+ def __call__(self, pred, shape_list):
267
+ pred = pred[:, 0, :, :]
268
+ segmentation = pred > self.thresh
269
+ boxes_batch = []
270
+ for batch_index in range(pred.shape[0]):
271
+ src_h, src_w, ratio_h, ratio_w = shape_list[batch_index]
272
+ if self.dilation_kernel is not None:
273
+ mask = cv2.dilate(
274
+ np.array(segmentation[batch_index]).astype(np.uint8),
275
+ self.dilation_kernel,
276
+ )
277
+ else:
278
+ mask = segmentation[batch_index]
279
+ boxes, _ = self._boxes_from_bitmap(pred[batch_index], mask, src_w, src_h)
280
+ boxes_batch.append(boxes)
281
+ return boxes_batch
282
+
283
+
284
+ # ============================================================================
285
+ # 3. Detection IoU Evaluator (inline from eval_det_iou.py)
286
+ # ============================================================================
287
+
288
+ _Rectangle = namedtuple("Rectangle", "xmin ymin xmax ymax")
289
+
290
+
291
+ def _get_intersection(pD, pG):
292
+ return Polygon(pD).intersection(Polygon(pG)).area
293
+
294
+
295
+ def _get_union(pD, pG):
296
+ return Polygon(pD).union(Polygon(pG)).area
297
+
298
+
299
+ def _get_iou(pD, pG):
300
+ return _get_intersection(pD, pG) / _get_union(pD, pG)
301
+
302
+
303
+ class _DetectionIoUEvaluator:
304
+ def __init__(self, iou_constraint=0.5, area_precision_constraint=0.5):
305
+ self.iou_constraint = iou_constraint
306
+ self.area_precision_constraint = area_precision_constraint
307
+
308
+ def evaluate_image(self, gt: List[dict], det: List[dict]) -> dict:
309
+ gt_pols = []
310
+ gt_dont_care = []
311
+
312
+ for n, g in enumerate(gt):
313
+ points = g.get("points", [])
314
+ if not points:
315
+ continue
316
+ try:
317
+ if not Polygon(points).is_valid:
318
+ continue
319
+ except Exception:
320
+ continue
321
+ gt_pols.append(points)
322
+ if g.get("ignore", False):
323
+ gt_dont_care.append(len(gt_pols) - 1)
324
+
325
+ det_pols = []
326
+ det_dont_care = []
327
+
328
+ for n, d in enumerate(det):
329
+ points = d.get("points", [])
330
+ if not points:
331
+ continue
332
+ try:
333
+ if not Polygon(points).is_valid:
334
+ continue
335
+ except Exception:
336
+ continue
337
+ det_pols.append(points)
338
+ if gt_dont_care:
339
+ for dc_idx in gt_dont_care:
340
+ dc_pol = gt_pols[dc_idx]
341
+ inter = _get_intersection(dc_pol, points)
342
+ pd_area = Polygon(points).area
343
+ prec = 0 if pd_area == 0 else inter / pd_area
344
+ if prec > self.area_precision_constraint:
345
+ det_dont_care.append(len(det_pols) - 1)
346
+ break
347
+
348
+ det_matched = 0
349
+
350
+ if gt_pols and det_pols:
351
+ iou_mat = np.empty([len(gt_pols), len(det_pols)])
352
+ for g_i, g_pts in enumerate(gt_pols):
353
+ for d_i, d_pts in enumerate(det_pols):
354
+ iou_mat[g_i, d_i] = _get_iou(d_pts, g_pts)
355
+
356
+ gt_matched = np.zeros(len(gt_pols), dtype=np.uint8)
357
+ det_matched_arr = np.zeros(len(det_pols), dtype=np.uint8)
358
+ for g_i in range(len(gt_pols)):
359
+ for d_i in range(len(det_pols)):
360
+ if (
361
+ gt_matched[g_i] == 0
362
+ and det_matched_arr[d_i] == 0
363
+ and g_i not in gt_dont_care
364
+ and d_i not in det_dont_care
365
+ ):
366
+ if iou_mat[g_i, d_i] > self.iou_constraint:
367
+ gt_matched[g_i] = 1
368
+ det_matched_arr[d_i] = 1
369
+ det_matched += 1
370
+
371
+ num_gt_care = len(gt_pols) - len(gt_dont_care)
372
+ num_det_care = len(det_pols) - len(det_dont_care)
373
+
374
+ return {
375
+ "gt_care": num_gt_care,
376
+ "det_care": num_det_care,
377
+ "det_matched": det_matched,
378
+ }
379
+
380
+ def combine_results(self, results: List[dict]) -> dict:
381
+ num_gt = sum(r["gt_care"] for r in results)
382
+ num_det = sum(r["det_care"] for r in results)
383
+ matched = sum(r["det_matched"] for r in results)
384
+
385
+ recall = 0 if num_gt == 0 else float(matched) / num_gt
386
+ precision = 0 if num_det == 0 else float(matched) / num_det
387
+ hmean = (
388
+ 0
389
+ if (precision + recall) == 0
390
+ else 2.0 * precision * recall / (precision + recall)
391
+ )
392
+
393
+ return {
394
+ "precision": precision,
395
+ "recall": recall,
396
+ "hmean": hmean,
397
+ }
398
+
399
+
400
+ # ============================================================================
401
+ # 4. Visualization
402
+ # ============================================================================
403
+
404
+ def draw_det_result(
405
+ img: np.ndarray,
406
+ det_boxes: np.ndarray,
407
+ gt_boxes: Optional[List[dict]] = None,
408
+ matched_pairs: Optional[List[Tuple[int, int]]] = None,
409
+ thickness: int = 2,
410
+ ) -> np.ndarray:
411
+ """Draw detection boxes on image with optional ground-truth and matching info.
412
+
413
+ Colors:
414
+ green = matched detection
415
+ blue = unmatched detection
416
+ red = unmatched ground truth (missed)
417
+ """
418
+ vis = img.copy()
419
+
420
+ if gt_boxes is not None and matched_pairs is not None:
421
+ gt_matched = set()
422
+ det_matched = set()
423
+ for p in matched_pairs:
424
+ gt_matched.add(p["gt"])
425
+ det_matched.add(p["det"])
426
+
427
+ for i, d in enumerate(det_boxes):
428
+ box = np.array(d, dtype=np.int32).reshape((-1, 1, 2))
429
+ color = (0, 255, 0) if i in det_matched else (255, 0, 0)
430
+ cv2.polylines(vis, [box], True, color, thickness)
431
+
432
+ for i, g in enumerate(gt_boxes):
433
+ if g.get("ignore", False):
434
+ continue
435
+ if i not in gt_matched:
436
+ pts = np.array(g["points"], dtype=np.int32).reshape((-1, 1, 2))
437
+ cv2.polylines(vis, [pts], True, (0, 0, 255), max(thickness, 3))
438
+ # Draw dashed effect by alternating segments
439
+ cx, cy = int(np.mean(pts[:, 0, 0])), int(np.mean(pts[:, 0, 1]))
440
+ cv2.putText(vis, "MISS", (cx, cy),
441
+ cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1)
442
+ else:
443
+ for box in det_boxes:
444
+ box = np.array(box, dtype=np.int32).reshape((-1, 1, 2))
445
+ cv2.polylines(vis, [box], True, (0, 255, 0), thickness)
446
+
447
+ return vis
448
+
449
+
450
+ # ============================================================================
451
+ # 5. Detection Engine
452
+ # ============================================================================
453
+
454
+ class PPOCRv6DetOnnx:
455
+
456
+ def __init__(
457
+ self,
458
+ det_onnx: str,
459
+ det_limit_side_len: int = 960,
460
+ det_db_thresh: float = 0.2,
461
+ det_db_box_thresh: float = 0.4,
462
+ det_db_unclip_ratio: float = 1.4,
463
+ det_max_candidates: int = 3000,
464
+ use_gpu: bool = False,
465
+ onnx_providers: Optional[List[str]] = None,
466
+ resize_mode: str = "letterbox",
467
+ ):
468
+ assert resize_mode in ("letterbox", "stretch"), f"invalid resize_mode: {resize_mode}"
469
+
470
+ self.session = ort.InferenceSession(det_onnx)
471
+ self.input_name = self.session.get_inputs()[0].name
472
+
473
+ # Detect fixed vs dynamic input dimensions
474
+ det_input = self.session.get_inputs()[0]
475
+ img_h = _get_dim_value(det_input.shape[2])
476
+ img_w = _get_dim_value(det_input.shape[3])
477
+
478
+ self._fixed_h = img_h if img_h > 0 else 0
479
+ self._fixed_w = img_w if img_w > 0 else 0
480
+
481
+ print(f"[PPOCRv6Det] ONNX input shape: {det_input.shape}, fixed_h={self._fixed_h}, fixed_w={self._fixed_w}, resize_mode={resize_mode}")
482
+
483
+ self._resize_mode = resize_mode
484
+
485
+ # Preprocessing: choose resize strategy
486
+ if self._fixed_h > 0 and self._fixed_w > 0:
487
+ # Fully fixed ONNX input — handled in _preprocess
488
+ self._resize_style = "fixed"
489
+ else:
490
+ # Dynamic or partially-fixed — ratio-preserving resize
491
+ self._resize = _DetResizeForTest(
492
+ limit_side_len=det_limit_side_len, limit_type="max"
493
+ )
494
+ self._resize_style = "dynamic"
495
+
496
+ self._normalize = _NormalizeImage(
497
+ mean=[0., 0., 0.],
498
+ std=[1.0, 1.0, 1.0],
499
+ scale=1.0,
500
+ )
501
+
502
+ self._to_chw = _ToCHWImage()
503
+ self._post = _DBPostProcess(
504
+ thresh=det_db_thresh,
505
+ box_thresh=det_db_box_thresh,
506
+ unclip_ratio=det_db_unclip_ratio,
507
+ max_candidates=det_max_candidates,
508
+ )
509
+
510
+ def _preprocess(self, img: np.ndarray):
511
+ src_h, src_w = img.shape[:2]
512
+ fixed_w = self._fixed_w
513
+ fixed_h = self._fixed_h
514
+
515
+ # Stretch mode: direct resize to fixed size (official PaddleOCR behavior)
516
+ if self._resize_mode == "stretch" and fixed_h > 0 and fixed_w > 0:
517
+ img_resized = cv2.resize(img, (fixed_w, fixed_h))
518
+ ratio_h = float(fixed_h) / src_h
519
+ ratio_w = float(fixed_w) / src_w
520
+ # Post-processing maps: origin = fm_coord / fm_dim * dest_dim
521
+ # For direct stretch, fm_dim corresponds uniformly to src_dim.
522
+ shape = np.array([src_h, src_w, ratio_h, ratio_w])
523
+
524
+ elif fixed_w > 0 or fixed_h > 0:
525
+ # Letterbox mode (default): ratio-preserving + pad to fixed size
526
+ ratios = []
527
+ if fixed_w > 0:
528
+ ratios.append(fixed_w / src_w)
529
+ if fixed_h > 0:
530
+ ratios.append(fixed_h / src_h)
531
+ ratio = min(ratios)
532
+ new_w = max(int(round(src_w * ratio / 32) * 32), 32)
533
+ new_h = max(int(round(src_h * ratio / 32) * 32), 32)
534
+ new_w = min(new_w, fixed_w) if fixed_w > 0 else new_w
535
+ new_h = min(new_h, fixed_h) if fixed_h > 0 else new_h
536
+
537
+ img_resized = cv2.resize(img, (new_w, new_h))
538
+ ratio_h = new_h / float(src_h)
539
+ ratio_w = new_w / float(src_w)
540
+
541
+ pad_h = max(0, fixed_h - new_h)
542
+ pad_w = max(0, fixed_w - new_w)
543
+ if pad_h > 0 or pad_w > 0:
544
+ img_resized = cv2.copyMakeBorder(
545
+ img_resized, 0, pad_h, 0, pad_w,
546
+ cv2.BORDER_CONSTANT, value=(0, 0, 0),
547
+ )
548
+
549
+ # Adjust shape for correct coordinate mapping after padding
550
+ adj_h = src_h * fixed_h / new_h if fixed_h > 0 else src_h
551
+ adj_w = src_w * fixed_w / new_w if fixed_w > 0 else src_w
552
+ shape = np.array([adj_h, adj_w, ratio_h, ratio_w])
553
+
554
+ else:
555
+ img_resized, shape = self._resize(img)
556
+
557
+ img_norm = self._normalize(img_resized)
558
+ img_chw = self._to_chw(img_norm)
559
+ tensor = np.expand_dims(img_chw.astype(np.float32), axis=0)
560
+ return tensor, shape
561
+
562
+ def _postprocess(self, output: np.ndarray, shape: np.ndarray):
563
+ shape_list = np.expand_dims(shape, axis=0)
564
+ boxes_batch = self._post(output, shape_list)
565
+ return boxes_batch[0]
566
+
567
+ def __call__(self, img: np.ndarray) -> np.ndarray:
568
+ """Detect text boxes. Returns (N, 4, 2) int32 array."""
569
+ tensor, shape = self._preprocess(img)
570
+ onnx_out = self.session.run(None, {self.input_name: tensor})
571
+ boxes = self._postprocess(onnx_out[0], shape)
572
+ return boxes
573
+
574
+ def predict_image(self, path: str) -> np.ndarray:
575
+ im = cv2.imread(path)
576
+ if im is None:
577
+ raise FileNotFoundError(f"Cannot read: {path}")
578
+ return self.__call__(im)
579
+
580
+
581
+ # ============================================================================
582
+ # 6. Evaluation
583
+ # ============================================================================
584
+
585
+ def evaluate(
586
+ det: PPOCRv6DetOnnx,
587
+ label_file: str,
588
+ dataset_root: str = "",
589
+ iou_constraint: float = 0.5,
590
+ verbose: bool = False,
591
+ ) -> dict:
592
+ """Evaluate detection against a PaddleOCR format label file.
593
+
594
+ Label format (one per line, tab-separated)::
595
+
596
+ rel/path/to/img.jpg<TAB>[{"transcription":"text_or_###","points":[[x,y]*4]}, ...]
597
+
598
+ "###" means ignored / don't-care region.
599
+
600
+ Returns:
601
+ dict: precision, recall, hmean, total_images, total_gt, total_det, det_matched, per_sample
602
+ """
603
+ samples = []
604
+ with open(label_file, "r", encoding="utf-8") as f:
605
+ for line in f:
606
+ line = line.strip()
607
+ if not line:
608
+ continue
609
+ parts = line.split("\t")
610
+ if len(parts) < 2:
611
+ continue
612
+ img_path = os.path.join(dataset_root, parts[0].strip())
613
+ try:
614
+ gt_label = json.loads(parts[1])
615
+ except json.JSONDecodeError:
616
+ continue
617
+ if not isinstance(gt_label, list):
618
+ continue
619
+ samples.append((img_path, gt_label))
620
+
621
+ if not samples:
622
+ print("[WARN] No samples found in label file.")
623
+ return {
624
+ "precision": 0, "recall": 0, "hmean": 0,
625
+ "total_images": 0, "total_gt": 0, "total_det": 0, "det_matched": 0,
626
+ "per_sample": [],
627
+ }
628
+
629
+ evaluator = _DetectionIoUEvaluator(iou_constraint=iou_constraint)
630
+ per_sample = []
631
+ total_gt = 0
632
+ total_det = 0
633
+ total_matched = 0
634
+
635
+ for idx, (img_path, gt_label) in enumerate(samples):
636
+ img = cv2.imread(img_path)
637
+ if img is None:
638
+ print(f"[WARN] Cannot read {img_path}, skipping.")
639
+ per_sample.append({
640
+ "image": img_path, "error": "cannot read",
641
+ "gt_care": 0, "det_care": 0, "det_matched": 0,
642
+ })
643
+ continue
644
+
645
+ # Run detection
646
+ det_boxes = det(img)
647
+
648
+ # Prepare GT
649
+ gt_info = []
650
+ for g in gt_label:
651
+ pts = g.get("points", [])
652
+ if not pts or len(pts) < 4:
653
+ continue
654
+ is_ignored = g.get("transcription", "") == "###"
655
+ gt_info.append({"points": pts, "ignore": is_ignored})
656
+
657
+ # Prepare DET
658
+ det_info = [{"points": d.tolist()} for d in det_boxes]
659
+
660
+ # Evaluate
661
+ result = evaluator.evaluate_image(gt_info, det_info)
662
+ total_gt += result["gt_care"]
663
+ total_det += result["det_care"]
664
+ total_matched += result["det_matched"]
665
+
666
+ per_sample.append({
667
+ "image": img_path,
668
+ "gt_care": result["gt_care"],
669
+ "det_care": result["det_care"],
670
+ "det_matched": result["det_matched"],
671
+ "det_boxes": [d.tolist() for d in det_boxes],
672
+ })
673
+
674
+ if verbose:
675
+ miss = result["gt_care"] - result["det_matched"]
676
+ extra = result["det_care"] - result["det_matched"]
677
+ parts = [
678
+ f"gt={result['gt_care']}",
679
+ f"det={result['det_care']}",
680
+ f"match={result['det_matched']}",
681
+ ]
682
+ if miss > 0:
683
+ parts.append(f"MISS={miss}")
684
+ if extra > 0:
685
+ parts.append(f"EXTRA={extra}")
686
+ print(f"[{os.path.basename(img_path)}] " + " ".join(parts))
687
+
688
+ metrics = evaluator.combine_results(
689
+ [r for r in per_sample if "error" not in r]
690
+ )
691
+
692
+ return {
693
+ "precision": round(metrics["precision"], 6),
694
+ "recall": round(metrics["recall"], 6),
695
+ "hmean": round(metrics["hmean"], 6),
696
+ "total_images": len(samples),
697
+ "total_gt": total_gt,
698
+ "total_det": total_det,
699
+ "det_matched": total_matched,
700
+ "per_sample": per_sample,
701
+ }
702
+
703
+
704
+ # ============================================================================
705
+ # 7. CLI
706
+ # ============================================================================
707
+
708
+ def main():
709
+ parser = argparse.ArgumentParser(
710
+ description="PP-OCRv6 Detection ONNX – inference & evaluation"
711
+ )
712
+ # Model
713
+ parser.add_argument("--det_onnx", type=str,
714
+ default="./axmodel/ax650/det_npu1.axmodel",
715
+ help="Path to detection ONNX model")
716
+ parser.add_argument("--limit_side_len", type=int, default=960)
717
+ parser.add_argument("--det_db_thresh", type=float, default=0.2)
718
+ parser.add_argument("--det_db_box_thresh", type=float, default=0.45)
719
+ parser.add_argument("--det_db_unclip_ratio", type=float, default=1.4)
720
+ parser.add_argument("--use_gpu", action="store_true", help="Enable GPU inference")
721
+ parser.add_argument("--resize_mode", type=str, default="letterbox",
722
+ choices=["letterbox", "stretch"],
723
+ help="Resize strategy for fixed-size ONNX: letterbox (keep ratio+pad) or stretch (direct resize)")
724
+
725
+ # Single image mode
726
+ parser.add_argument("--image", type=str, default=None, help="Single image path")
727
+
728
+ # Evaluation mode
729
+ parser.add_argument("--label_file", type=str,
730
+ default='dataset/ocr_det_dataset_examples/val.txt',
731
+ help="Label file (image_path<TAB>json_label per line)")
732
+ parser.add_argument("--dataset_root", type=str,
733
+ default="dataset/ocr_det_dataset_examples",
734
+ help="Prefix directory for image paths in label file")
735
+
736
+ # Common
737
+ parser.add_argument("--visualize", action="store_true", help="Draw boxes on image")
738
+ parser.add_argument("--output", type=str, default=None,
739
+ help="Save visualized image (implies --visualize)")
740
+ parser.add_argument("--verbose", action="store_true", help="Print per-image metrics")
741
+ parser.add_argument("--output_json", type=str, default=None,
742
+ help="Save results to JSON file")
743
+
744
+ args = parser.parse_args()
745
+
746
+ det = PPOCRv6DetOnnx(
747
+ det_onnx=args.det_onnx,
748
+ det_limit_side_len=args.limit_side_len,
749
+ det_db_thresh=args.det_db_thresh,
750
+ det_db_box_thresh=args.det_db_box_thresh,
751
+ det_db_unclip_ratio=args.det_db_unclip_ratio,
752
+ use_gpu=args.use_gpu,
753
+ resize_mode=args.resize_mode,
754
+ )
755
+
756
+ # --- Single image mode ---
757
+ if args.image and not args.label_file:
758
+ img = cv2.imread(args.image)
759
+ if img is None:
760
+ raise FileNotFoundError(f"Cannot read: {args.image}")
761
+ boxes = det(img)
762
+ print(f"Detected {len(boxes)} text boxes:")
763
+ for i, box in enumerate(boxes):
764
+ print(f" [{i}] {box.tolist()}")
765
+
766
+ do_viz = args.visualize or args.output
767
+ if do_viz:
768
+ vis = draw_det_result(img, boxes)
769
+ out_path = args.output or "det_result.jpg"
770
+ cv2.imwrite(out_path, vis)
771
+ print(f"Visualization saved to: {out_path}")
772
+
773
+ if args.output_json:
774
+ with open(args.output_json, "w") as f:
775
+ json.dump(
776
+ {"image": args.image, "boxes": [b.tolist() for b in boxes]},
777
+ f, indent=2,
778
+ )
779
+ print(f"Results saved to: {args.output_json}")
780
+ return
781
+
782
+ # --- Evaluation mode ---
783
+ if args.label_file:
784
+ metrics = evaluate(
785
+ det, args.label_file,
786
+ dataset_root=args.dataset_root,
787
+ verbose=args.verbose,
788
+ )
789
+
790
+ print()
791
+ print("=" * 60)
792
+ print("Evaluation Results")
793
+ print("=" * 60)
794
+ print(f" Images: {metrics['total_images']}")
795
+ print(f" GT boxes: {metrics['total_gt']}")
796
+ print(f" DET boxes: {metrics['total_det']}")
797
+ print(f" Matched: {metrics['det_matched']}")
798
+ print(f" Precision: {metrics['precision']:.4f} ({metrics['precision']*100:.2f}%)")
799
+ print(f" Recall: {metrics['recall']:.4f} ({metrics['recall']*100:.2f}%)")
800
+ print(f" Hmean (F1): {metrics['hmean']:.4f}")
801
+ print("=" * 60)
802
+
803
+ # Visualization for eval mode
804
+ do_viz = args.visualize or args.output
805
+ if do_viz:
806
+ out_dir = args.output if args.output else "det_eval_vis"
807
+ os.makedirs(out_dir, exist_ok=True)
808
+ for i, smp in enumerate(metrics["per_sample"]):
809
+ img = cv2.imread(smp["image"])
810
+ if img is None:
811
+ continue
812
+ # Load GT boxes with matching info
813
+ with open(args.label_file, "r") as f:
814
+ lines = f.readlines()
815
+ gt_label = []
816
+ for line in lines:
817
+ line = line.strip()
818
+ if not line:
819
+ continue
820
+ parts = line.split("\t")
821
+ if len(parts) < 2:
822
+ continue
823
+ if os.path.join(args.dataset_root, parts[0].strip()) == smp["image"]:
824
+ gt_label = json.loads(parts[1])
825
+ break
826
+
827
+ det_boxes = np.array(smp.get("det_boxes", []))
828
+ # Simple matching for visualization (re-run evaluate_image)
829
+ gt_info = []
830
+ for g in gt_label:
831
+ pts = g.get("points", [])
832
+ if not pts or len(pts) < 4:
833
+ continue
834
+ gt_info.append({
835
+ "points": pts,
836
+ "ignore": g.get("transcription", "") == "###",
837
+ })
838
+ det_info = [{"points": d} for d in det_boxes.tolist()]
839
+
840
+ # Compute matches for coloring
841
+ matched_pairs = _compute_matched_pairs(gt_info, det_info)
842
+ vis = draw_det_result(img, det_boxes, gt_info, matched_pairs)
843
+ fname = os.path.basename(smp["image"])
844
+ cv2.imwrite(os.path.join(out_dir, fname), vis)
845
+
846
+ print(f"Visualization saved to: {out_dir}/")
847
+
848
+ if args.output_json:
849
+ out = {k: v for k, v in metrics.items() if k != "per_sample"}
850
+ out["per_sample"] = metrics["per_sample"]
851
+ with open(args.output_json, "w") as f:
852
+ json.dump(out, f, indent=2)
853
+ print(f"Results saved to: {args.output_json}")
854
+ return
855
+
856
+ parser.error("Either --image or --label_file must be provided.")
857
+
858
+
859
+ def _compute_matched_pairs(gt_info, det_info, iou_thr=0.5):
860
+ """Compute matched GT-det pairs for visualization coloring."""
861
+ pairs = []
862
+ gt_pols = [g["points"] for g in gt_info if not g.get("ignore")]
863
+ det_pols = [d["points"] for d in det_info]
864
+ if not gt_pols or not det_pols:
865
+ return pairs
866
+
867
+ iou_mat = np.empty([len(gt_pols), len(det_pols)])
868
+ for g_i, g_pts in enumerate(gt_pols):
869
+ for d_i, d_pts in enumerate(det_pols):
870
+ try:
871
+ int_area = _get_intersection(d_pts, g_pts)
872
+ union_area = _get_union(d_pts, g_pts)
873
+ iou_mat[g_i, d_i] = int_area / union_area if union_area > 0 else 0
874
+ except Exception:
875
+ iou_mat[g_i, d_i] = 0
876
+
877
+ gt_used = set()
878
+ det_used = set()
879
+ # Greedy matching by descending IoU
880
+ flat = []
881
+ for g_i in range(len(gt_pols)):
882
+ for d_i in range(len(det_pols)):
883
+ flat.append((iou_mat[g_i, d_i], g_i, d_i))
884
+ flat.sort(key=lambda x: x[0], reverse=True)
885
+ for iou, g_i, d_i in flat:
886
+ if iou > iou_thr and g_i not in gt_used and d_i not in det_used:
887
+ pairs.append({"gt": g_i, "det": d_i})
888
+ gt_used.add(g_i)
889
+ det_used.add(d_i)
890
+ return pairs
891
+
892
+
893
+ if __name__ == "__main__":
894
+ main()
run_det_onnx.py ADDED
@@ -0,0 +1,913 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ PP-OCRv6 Detection ONNX Inference & Evaluation (standalone, zero Paddle dependency)
4
+
5
+ Dependencies:
6
+ numpy, opencv-python, onnxruntime, pyyaml, shapely, pyclipper
7
+
8
+ Data format (same as PaddleOCR official):
9
+ Label file: image_path<TAB>json_label
10
+ json_label: [{"transcription": "text_or_###", "points": [[x,y]*4]}, ...]
11
+
12
+ Usage:
13
+ # Single image inference
14
+ python ppocrv6_det_onnx.py --det_onnx det.onnx --image test.jpg --visualize
15
+
16
+ # Batch evaluation
17
+ python ppocrv6_det_onnx.py --det_onnx det.onnx \\
18
+ --label_file val.txt --dataset_root ./ocr_det_dataset_examples \\
19
+ --visualize --output_json result.json --verbose
20
+ """
21
+
22
+ import argparse
23
+ import json
24
+ import math
25
+ import os
26
+ from collections import namedtuple
27
+ from typing import List, Optional, Tuple
28
+
29
+ import cv2
30
+ import numpy as np
31
+ import onnxruntime as ort
32
+ import yaml
33
+ from shapely.geometry import Polygon
34
+ import pyclipper
35
+
36
+
37
+ def _get_dim_value(dim):
38
+ """Extract integer value from an ONNX Runtime dimension, returning 0 for dynamic dims."""
39
+ if dim is None:
40
+ return 0
41
+ if isinstance(dim, str):
42
+ return 0
43
+ if hasattr(dim, 'dim_value'):
44
+ return int(dim.dim_value) if dim.dim_value else 0
45
+ if hasattr(dim, 'dim_param'):
46
+ # named dim like "batch_size" → dynamic
47
+ return 0
48
+ try:
49
+ v = int(dim)
50
+ return v
51
+ except (TypeError, ValueError):
52
+ return 0
53
+
54
+
55
+ # ============================================================================
56
+ # 1. Detection Preprocessing
57
+ # ============================================================================
58
+
59
+ class _DetResizeForTest:
60
+ def __init__(self, limit_side_len=960, limit_type="max", max_side_limit=4000,
61
+ image_shape=None, keep_ratio=False):
62
+ self.max_side_limit = max_side_limit
63
+ # resize_type=0: limit_side_len (dynamic input)
64
+ # resize_type=1: image_shape (fixed ONNX input)
65
+ if image_shape is not None:
66
+ self.resize_type = 1
67
+ self.image_shape = image_shape
68
+ self.keep_ratio = keep_ratio
69
+ else:
70
+ self.resize_type = 0
71
+ self.limit_side_len = limit_side_len
72
+ self.limit_type = limit_type
73
+
74
+ def _image_padding(self, im, value=0):
75
+ h, w, c = im.shape
76
+ im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value
77
+ im_pad[:h, :w, :] = im
78
+ return im_pad
79
+
80
+ def _resize_image_type0(self, img):
81
+ h, w, _ = img.shape
82
+ limit_side_len = self.limit_side_len
83
+
84
+ if self.limit_type == "max":
85
+ if max(h, w) > limit_side_len:
86
+ ratio = float(limit_side_len) / max(h, w)
87
+ else:
88
+ ratio = 1.0
89
+ elif self.limit_type == "min":
90
+ if min(h, w) < limit_side_len:
91
+ ratio = float(limit_side_len) / min(h, w)
92
+ else:
93
+ ratio = 1.0
94
+ elif self.limit_type == "resize_long":
95
+ ratio = float(limit_side_len) / max(h, w)
96
+ else:
97
+ raise ValueError(f"not support limit_type: {self.limit_type}")
98
+
99
+ resize_h = int(h * ratio)
100
+ resize_w = int(w * ratio)
101
+ if max(resize_h, resize_w) > self.max_side_limit:
102
+ ratio = float(self.max_side_limit) / max(resize_h, resize_w)
103
+ resize_h, resize_w = int(resize_h * ratio), int(resize_w * ratio)
104
+
105
+ resize_h = max(int(round(resize_h / 32) * 32), 32)
106
+ resize_w = max(int(round(resize_w / 32) * 32), 32)
107
+
108
+ if int(resize_w) <= 0 or int(resize_h) <= 0:
109
+ return None, (None, None)
110
+ img = cv2.resize(img, (int(resize_w), int(resize_h)))
111
+ ratio_h = resize_h / float(h)
112
+ ratio_w = resize_w / float(w)
113
+ return img, [ratio_h, ratio_w]
114
+
115
+ def _resize_image_type1(self, img):
116
+ """Direct resize to fixed [H, W]. Used when ONNX has fixed input dimensions."""
117
+ resize_h, resize_w = self.image_shape
118
+ ori_h, ori_w = img.shape[:2]
119
+ if self.keep_ratio:
120
+ resize_w = ori_w * resize_h / ori_h
121
+ N = math.ceil(resize_w / 32)
122
+ resize_w = N * 32
123
+ ratio_h = float(resize_h) / ori_h
124
+ ratio_w = float(resize_w) / ori_w
125
+ img = cv2.resize(img, (int(resize_w), int(resize_h)))
126
+ return img, [ratio_h, ratio_w]
127
+
128
+ def __call__(self, img):
129
+ src_h, src_w = img.shape[:2]
130
+ if sum([src_h, src_w]) < 64:
131
+ img = self._image_padding(img)
132
+ if self.resize_type == 1:
133
+ img, [ratio_h, ratio_w] = self._resize_image_type1(img)
134
+ else:
135
+ img, [ratio_h, ratio_w] = self._resize_image_type0(img)
136
+ shape = np.array([src_h, src_w, ratio_h, ratio_w])
137
+ return img, shape
138
+
139
+
140
+ class _NormalizeImage:
141
+ def __init__(self, mean, std, scale=1.0 / 255.0, order="hwc"):
142
+ self.scale = np.float32(scale)
143
+ shape = (1, 1, 3) if order == "hwc" else (3, 1, 1)
144
+ self.mean = np.array(mean, dtype=np.float32).reshape(shape)
145
+ self.std = np.array(std, dtype=np.float32).reshape(shape)
146
+
147
+ def __call__(self, img):
148
+ return (img.astype("float32") * self.scale - self.mean) / self.std
149
+
150
+
151
+ class _ToCHWImage:
152
+ def __call__(self, img):
153
+ return img.transpose((2, 0, 1))
154
+
155
+
156
+ # ============================================================================
157
+ # 2. Detection Postprocessing (DB)
158
+ # ============================================================================
159
+
160
+ class _DBPostProcess:
161
+ def __init__(
162
+ self,
163
+ thresh=0.3,
164
+ box_thresh=0.7,
165
+ max_candidates=1000,
166
+ unclip_ratio=2.0,
167
+ use_dilation=False,
168
+ score_mode="fast",
169
+ ):
170
+ self.thresh = thresh
171
+ self.box_thresh = box_thresh
172
+ self.max_candidates = max_candidates
173
+ self.unclip_ratio = unclip_ratio
174
+ self.min_size = 3
175
+ self.score_mode = score_mode
176
+ assert score_mode in ("slow", "fast")
177
+ self.dilation_kernel = None if not use_dilation else np.array([[1, 1], [1, 1]])
178
+
179
+ def _unclip(self, box, unclip_ratio):
180
+ poly = Polygon(box)
181
+ distance = poly.area * unclip_ratio / poly.length
182
+ offset = pyclipper.PyclipperOffset()
183
+ offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
184
+ expanded = offset.Execute(distance)
185
+ return expanded
186
+
187
+ def _get_mini_boxes(self, contour):
188
+ bounding_box = cv2.minAreaRect(contour)
189
+ points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
190
+ i1, i2, i3, i4 = 0, 1, 2, 3
191
+ if points[1][1] > points[0][1]:
192
+ i1, i4 = 0, 1
193
+ else:
194
+ i1, i4 = 1, 0
195
+ if points[3][1] > points[2][1]:
196
+ i2, i3 = 2, 3
197
+ else:
198
+ i2, i3 = 3, 2
199
+ box = [points[i1], points[i2], points[i3], points[i4]]
200
+ return box, min(bounding_box[1])
201
+
202
+ def _box_score_fast(self, bitmap, _box):
203
+ h, w = bitmap.shape[:2]
204
+ box = _box.copy()
205
+ xmin = np.clip(np.floor(box[:, 0].min()).astype("int32"), 0, w - 1)
206
+ xmax = np.clip(np.ceil(box[:, 0].max()).astype("int32"), 0, w - 1)
207
+ ymin = np.clip(np.floor(box[:, 1].min()).astype("int32"), 0, h - 1)
208
+ ymax = np.clip(np.ceil(box[:, 1].max()).astype("int32"), 0, h - 1)
209
+ mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
210
+ box[:, 0] = box[:, 0] - xmin
211
+ box[:, 1] = box[:, 1] - ymin
212
+ cv2.fillPoly(mask, box.reshape(1, -1, 2).astype("int32"), 1)
213
+ return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]
214
+
215
+ def _boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
216
+ bitmap = _bitmap
217
+ height, width = bitmap.shape
218
+ outs = cv2.findContours(
219
+ (bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
220
+ )
221
+ if len(outs) == 3:
222
+ _, contours, _ = outs
223
+ else:
224
+ contours, _ = outs
225
+ num_contours = min(len(contours), self.max_candidates)
226
+ boxes, scores = [], []
227
+ for index in range(num_contours):
228
+ contour = contours[index]
229
+ points, sside = self._get_mini_boxes(contour)
230
+ if sside < self.min_size:
231
+ continue
232
+ points = np.array(points)
233
+ if self.score_mode == "fast":
234
+ score = self._box_score_fast(pred, points.reshape(-1, 2))
235
+ else:
236
+ score = self._box_score_slow(pred, contour)
237
+ if self.box_thresh > score:
238
+ continue
239
+ box = self._unclip(points, self.unclip_ratio)
240
+ if len(box) > 1:
241
+ continue
242
+ box = np.array(box).reshape(-1, 1, 2)
243
+ box, sside = self._get_mini_boxes(box)
244
+ if sside < self.min_size + 2:
245
+ continue
246
+ box = np.array(box)
247
+ box[:, 0] = np.clip(np.round(box[:, 0] / width * dest_width), 0, dest_width)
248
+ box[:, 1] = np.clip(np.round(box[:, 1] / height * dest_height), 0, dest_height)
249
+ boxes.append(box.astype("int32"))
250
+ scores.append(score)
251
+ return np.array(boxes, dtype="int32"), scores
252
+
253
+ def _box_score_slow(self, bitmap, contour):
254
+ h, w = bitmap.shape[:2]
255
+ contour = contour.copy().reshape((-1, 2))
256
+ xmin = np.clip(np.min(contour[:, 0]), 0, w - 1)
257
+ xmax = np.clip(np.max(contour[:, 0]), 0, w - 1)
258
+ ymin = np.clip(np.min(contour[:, 1]), 0, h - 1)
259
+ ymax = np.clip(np.max(contour[:, 1]), 0, h - 1)
260
+ mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
261
+ contour[:, 0] = contour[:, 0] - xmin
262
+ contour[:, 1] = contour[:, 1] - ymin
263
+ cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype("int32"), 1)
264
+ return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]
265
+
266
+ def __call__(self, pred, shape_list):
267
+ pred = pred[:, 0, :, :]
268
+ segmentation = pred > self.thresh
269
+ boxes_batch = []
270
+ for batch_index in range(pred.shape[0]):
271
+ src_h, src_w, ratio_h, ratio_w = shape_list[batch_index]
272
+ if self.dilation_kernel is not None:
273
+ mask = cv2.dilate(
274
+ np.array(segmentation[batch_index]).astype(np.uint8),
275
+ self.dilation_kernel,
276
+ )
277
+ else:
278
+ mask = segmentation[batch_index]
279
+ boxes, _ = self._boxes_from_bitmap(pred[batch_index], mask, src_w, src_h)
280
+ boxes_batch.append(boxes)
281
+ return boxes_batch
282
+
283
+
284
+ # ============================================================================
285
+ # 3. Detection IoU Evaluator (inline from eval_det_iou.py)
286
+ # ============================================================================
287
+
288
+ _Rectangle = namedtuple("Rectangle", "xmin ymin xmax ymax")
289
+
290
+
291
+ def _get_intersection(pD, pG):
292
+ return Polygon(pD).intersection(Polygon(pG)).area
293
+
294
+
295
+ def _get_union(pD, pG):
296
+ return Polygon(pD).union(Polygon(pG)).area
297
+
298
+
299
+ def _get_iou(pD, pG):
300
+ return _get_intersection(pD, pG) / _get_union(pD, pG)
301
+
302
+
303
+ class _DetectionIoUEvaluator:
304
+ def __init__(self, iou_constraint=0.5, area_precision_constraint=0.5):
305
+ self.iou_constraint = iou_constraint
306
+ self.area_precision_constraint = area_precision_constraint
307
+
308
+ def evaluate_image(self, gt: List[dict], det: List[dict]) -> dict:
309
+ gt_pols = []
310
+ gt_dont_care = []
311
+
312
+ for n, g in enumerate(gt):
313
+ points = g.get("points", [])
314
+ if not points:
315
+ continue
316
+ try:
317
+ if not Polygon(points).is_valid:
318
+ continue
319
+ except Exception:
320
+ continue
321
+ gt_pols.append(points)
322
+ if g.get("ignore", False):
323
+ gt_dont_care.append(len(gt_pols) - 1)
324
+
325
+ det_pols = []
326
+ det_dont_care = []
327
+
328
+ for n, d in enumerate(det):
329
+ points = d.get("points", [])
330
+ if not points:
331
+ continue
332
+ try:
333
+ if not Polygon(points).is_valid:
334
+ continue
335
+ except Exception:
336
+ continue
337
+ det_pols.append(points)
338
+ if gt_dont_care:
339
+ for dc_idx in gt_dont_care:
340
+ dc_pol = gt_pols[dc_idx]
341
+ inter = _get_intersection(dc_pol, points)
342
+ pd_area = Polygon(points).area
343
+ prec = 0 if pd_area == 0 else inter / pd_area
344
+ if prec > self.area_precision_constraint:
345
+ det_dont_care.append(len(det_pols) - 1)
346
+ break
347
+
348
+ det_matched = 0
349
+
350
+ if gt_pols and det_pols:
351
+ iou_mat = np.empty([len(gt_pols), len(det_pols)])
352
+ for g_i, g_pts in enumerate(gt_pols):
353
+ for d_i, d_pts in enumerate(det_pols):
354
+ iou_mat[g_i, d_i] = _get_iou(d_pts, g_pts)
355
+
356
+ gt_matched = np.zeros(len(gt_pols), dtype=np.uint8)
357
+ det_matched_arr = np.zeros(len(det_pols), dtype=np.uint8)
358
+ for g_i in range(len(gt_pols)):
359
+ for d_i in range(len(det_pols)):
360
+ if (
361
+ gt_matched[g_i] == 0
362
+ and det_matched_arr[d_i] == 0
363
+ and g_i not in gt_dont_care
364
+ and d_i not in det_dont_care
365
+ ):
366
+ if iou_mat[g_i, d_i] > self.iou_constraint:
367
+ gt_matched[g_i] = 1
368
+ det_matched_arr[d_i] = 1
369
+ det_matched += 1
370
+
371
+ num_gt_care = len(gt_pols) - len(gt_dont_care)
372
+ num_det_care = len(det_pols) - len(det_dont_care)
373
+
374
+ return {
375
+ "gt_care": num_gt_care,
376
+ "det_care": num_det_care,
377
+ "det_matched": det_matched,
378
+ }
379
+
380
+ def combine_results(self, results: List[dict]) -> dict:
381
+ num_gt = sum(r["gt_care"] for r in results)
382
+ num_det = sum(r["det_care"] for r in results)
383
+ matched = sum(r["det_matched"] for r in results)
384
+
385
+ recall = 0 if num_gt == 0 else float(matched) / num_gt
386
+ precision = 0 if num_det == 0 else float(matched) / num_det
387
+ hmean = (
388
+ 0
389
+ if (precision + recall) == 0
390
+ else 2.0 * precision * recall / (precision + recall)
391
+ )
392
+
393
+ return {
394
+ "precision": precision,
395
+ "recall": recall,
396
+ "hmean": hmean,
397
+ }
398
+
399
+
400
+ # ============================================================================
401
+ # 4. Visualization
402
+ # ============================================================================
403
+
404
+ def draw_det_result(
405
+ img: np.ndarray,
406
+ det_boxes: np.ndarray,
407
+ gt_boxes: Optional[List[dict]] = None,
408
+ matched_pairs: Optional[List[Tuple[int, int]]] = None,
409
+ thickness: int = 2,
410
+ ) -> np.ndarray:
411
+ """Draw detection boxes on image with optional ground-truth and matching info.
412
+
413
+ Colors:
414
+ green = matched detection
415
+ blue = unmatched detection
416
+ red = unmatched ground truth (missed)
417
+ """
418
+ vis = img.copy()
419
+
420
+ if gt_boxes is not None and matched_pairs is not None:
421
+ gt_matched = set()
422
+ det_matched = set()
423
+ for p in matched_pairs:
424
+ gt_matched.add(p["gt"])
425
+ det_matched.add(p["det"])
426
+
427
+ for i, d in enumerate(det_boxes):
428
+ box = np.array(d, dtype=np.int32).reshape((-1, 1, 2))
429
+ color = (0, 255, 0) if i in det_matched else (255, 0, 0)
430
+ cv2.polylines(vis, [box], True, color, thickness)
431
+
432
+ for i, g in enumerate(gt_boxes):
433
+ if g.get("ignore", False):
434
+ continue
435
+ if i not in gt_matched:
436
+ pts = np.array(g["points"], dtype=np.int32).reshape((-1, 1, 2))
437
+ cv2.polylines(vis, [pts], True, (0, 0, 255), max(thickness, 3))
438
+ # Draw dashed effect by alternating segments
439
+ cx, cy = int(np.mean(pts[:, 0, 0])), int(np.mean(pts[:, 0, 1]))
440
+ cv2.putText(vis, "MISS", (cx, cy),
441
+ cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1)
442
+ else:
443
+ for box in det_boxes:
444
+ box = np.array(box, dtype=np.int32).reshape((-1, 1, 2))
445
+ cv2.polylines(vis, [box], True, (0, 255, 0), thickness)
446
+
447
+ return vis
448
+
449
+
450
+ # ============================================================================
451
+ # 5. Detection Engine
452
+ # ============================================================================
453
+
454
+ class PPOCRv6DetOnnx:
455
+
456
+ def __init__(
457
+ self,
458
+ det_onnx: str,
459
+ det_limit_side_len: int = 960,
460
+ det_db_thresh: float = 0.2,
461
+ det_db_box_thresh: float = 0.4,
462
+ det_db_unclip_ratio: float = 1.4,
463
+ det_max_candidates: int = 3000,
464
+ use_gpu: bool = False,
465
+ onnx_providers: Optional[List[str]] = None,
466
+ resize_mode: str = "letterbox",
467
+ ):
468
+ assert resize_mode in ("letterbox", "stretch"), f"invalid resize_mode: {resize_mode}"
469
+ # ONNX session
470
+ if onnx_providers is None:
471
+ onnx_providers = (
472
+ ["CUDAExecutionProvider", "CPUExecutionProvider"]
473
+ if use_gpu
474
+ else ["CPUExecutionProvider"]
475
+ )
476
+ sess_options = ort.SessionOptions()
477
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
478
+ self.session = ort.InferenceSession(
479
+ det_onnx, sess_options=sess_options, providers=onnx_providers
480
+ )
481
+ self.input_name = self.session.get_inputs()[0].name
482
+
483
+ # Detect fixed vs dynamic input dimensions
484
+ det_input = self.session.get_inputs()[0]
485
+ img_h = _get_dim_value(det_input.shape[2])
486
+ img_w = _get_dim_value(det_input.shape[3])
487
+
488
+ # Fallback: try onnx.load() if ORT returned ambiguous dims
489
+ if img_h == 0 and img_w == 0:
490
+ try:
491
+ import onnx
492
+ m = onnx.load(det_onnx)
493
+ inp = m.graph.input[0]
494
+ dims = inp.type.tensor_type.shape.dim
495
+ img_h = dims[2].dim_value if len(dims) > 2 else 0
496
+ img_w = dims[3].dim_value if len(dims) > 3 else 0
497
+ except Exception:
498
+ pass
499
+
500
+ self._fixed_h = img_h if img_h > 0 else 0
501
+ self._fixed_w = img_w if img_w > 0 else 0
502
+
503
+ print(f"[PPOCRv6Det] ONNX input shape: {det_input.shape}, fixed_h={self._fixed_h}, fixed_w={self._fixed_w}, resize_mode={resize_mode}")
504
+
505
+ self._resize_mode = resize_mode
506
+
507
+ # Preprocessing: choose resize strategy
508
+ if self._fixed_h > 0 and self._fixed_w > 0:
509
+ # Fully fixed ONNX input — handled in _preprocess
510
+ self._resize_style = "fixed"
511
+ else:
512
+ # Dynamic or partially-fixed — ratio-preserving resize
513
+ self._resize = _DetResizeForTest(
514
+ limit_side_len=det_limit_side_len, limit_type="max"
515
+ )
516
+ self._resize_style = "dynamic"
517
+ self._normalize = _NormalizeImage(
518
+ mean=[0.485, 0.456, 0.406],
519
+ std=[0.229, 0.224, 0.225],
520
+ )
521
+ self._to_chw = _ToCHWImage()
522
+ self._post = _DBPostProcess(
523
+ thresh=det_db_thresh,
524
+ box_thresh=det_db_box_thresh,
525
+ unclip_ratio=det_db_unclip_ratio,
526
+ max_candidates=det_max_candidates,
527
+ )
528
+
529
+ def _preprocess(self, img: np.ndarray):
530
+ src_h, src_w = img.shape[:2]
531
+ fixed_w = self._fixed_w
532
+ fixed_h = self._fixed_h
533
+
534
+ # Stretch mode: direct resize to fixed size (official PaddleOCR behavior)
535
+ if self._resize_mode == "stretch" and fixed_h > 0 and fixed_w > 0:
536
+ img_resized = cv2.resize(img, (fixed_w, fixed_h))
537
+ ratio_h = float(fixed_h) / src_h
538
+ ratio_w = float(fixed_w) / src_w
539
+ # Post-processing maps: origin = fm_coord / fm_dim * dest_dim
540
+ # For direct stretch, fm_dim corresponds uniformly to src_dim.
541
+ shape = np.array([src_h, src_w, ratio_h, ratio_w])
542
+
543
+ elif fixed_w > 0 or fixed_h > 0:
544
+ # Letterbox mode (default): ratio-preserving + pad to fixed size
545
+ ratios = []
546
+ if fixed_w > 0:
547
+ ratios.append(fixed_w / src_w)
548
+ if fixed_h > 0:
549
+ ratios.append(fixed_h / src_h)
550
+ ratio = min(ratios)
551
+ new_w = max(int(round(src_w * ratio / 32) * 32), 32)
552
+ new_h = max(int(round(src_h * ratio / 32) * 32), 32)
553
+ new_w = min(new_w, fixed_w) if fixed_w > 0 else new_w
554
+ new_h = min(new_h, fixed_h) if fixed_h > 0 else new_h
555
+
556
+ img_resized = cv2.resize(img, (new_w, new_h))
557
+ ratio_h = new_h / float(src_h)
558
+ ratio_w = new_w / float(src_w)
559
+
560
+ pad_h = max(0, fixed_h - new_h)
561
+ pad_w = max(0, fixed_w - new_w)
562
+ if pad_h > 0 or pad_w > 0:
563
+ img_resized = cv2.copyMakeBorder(
564
+ img_resized, 0, pad_h, 0, pad_w,
565
+ cv2.BORDER_CONSTANT, value=(0, 0, 0),
566
+ )
567
+
568
+ # Adjust shape for correct coordinate mapping after padding
569
+ adj_h = src_h * fixed_h / new_h if fixed_h > 0 else src_h
570
+ adj_w = src_w * fixed_w / new_w if fixed_w > 0 else src_w
571
+ shape = np.array([adj_h, adj_w, ratio_h, ratio_w])
572
+
573
+ else:
574
+ img_resized, shape = self._resize(img)
575
+
576
+ img_norm = self._normalize(img_resized)
577
+ img_chw = self._to_chw(img_norm)
578
+ tensor = np.expand_dims(img_chw.astype(np.float32), axis=0)
579
+ return tensor, shape
580
+
581
+ def _postprocess(self, output: np.ndarray, shape: np.ndarray):
582
+ shape_list = np.expand_dims(shape, axis=0)
583
+ boxes_batch = self._post(output, shape_list)
584
+ return boxes_batch[0]
585
+
586
+ def __call__(self, img: np.ndarray) -> np.ndarray:
587
+ """Detect text boxes. Returns (N, 4, 2) int32 array."""
588
+ tensor, shape = self._preprocess(img)
589
+ onnx_out = self.session.run(None, {self.input_name: tensor})
590
+ boxes = self._postprocess(onnx_out[0], shape)
591
+ return boxes
592
+
593
+ def predict_image(self, path: str) -> np.ndarray:
594
+ im = cv2.imread(path)
595
+ if im is None:
596
+ raise FileNotFoundError(f"Cannot read: {path}")
597
+ return self.__call__(im)
598
+
599
+
600
+ # ============================================================================
601
+ # 6. Evaluation
602
+ # ============================================================================
603
+
604
+ def evaluate(
605
+ det: PPOCRv6DetOnnx,
606
+ label_file: str,
607
+ dataset_root: str = "",
608
+ iou_constraint: float = 0.5,
609
+ verbose: bool = False,
610
+ ) -> dict:
611
+ """Evaluate detection against a PaddleOCR format label file.
612
+
613
+ Label format (one per line, tab-separated)::
614
+
615
+ rel/path/to/img.jpg<TAB>[{"transcription":"text_or_###","points":[[x,y]*4]}, ...]
616
+
617
+ "###" means ignored / don't-care region.
618
+
619
+ Returns:
620
+ dict: precision, recall, hmean, total_images, total_gt, total_det, det_matched, per_sample
621
+ """
622
+ samples = []
623
+ with open(label_file, "r", encoding="utf-8") as f:
624
+ for line in f:
625
+ line = line.strip()
626
+ if not line:
627
+ continue
628
+ parts = line.split("\t")
629
+ if len(parts) < 2:
630
+ continue
631
+ img_path = os.path.join(dataset_root, parts[0].strip())
632
+ try:
633
+ gt_label = json.loads(parts[1])
634
+ except json.JSONDecodeError:
635
+ continue
636
+ if not isinstance(gt_label, list):
637
+ continue
638
+ samples.append((img_path, gt_label))
639
+
640
+ if not samples:
641
+ print("[WARN] No samples found in label file.")
642
+ return {
643
+ "precision": 0, "recall": 0, "hmean": 0,
644
+ "total_images": 0, "total_gt": 0, "total_det": 0, "det_matched": 0,
645
+ "per_sample": [],
646
+ }
647
+
648
+ evaluator = _DetectionIoUEvaluator(iou_constraint=iou_constraint)
649
+ per_sample = []
650
+ total_gt = 0
651
+ total_det = 0
652
+ total_matched = 0
653
+
654
+ for idx, (img_path, gt_label) in enumerate(samples):
655
+ img = cv2.imread(img_path)
656
+ if img is None:
657
+ print(f"[WARN] Cannot read {img_path}, skipping.")
658
+ per_sample.append({
659
+ "image": img_path, "error": "cannot read",
660
+ "gt_care": 0, "det_care": 0, "det_matched": 0,
661
+ })
662
+ continue
663
+
664
+ # Run detection
665
+ det_boxes = det(img)
666
+
667
+ # Prepare GT
668
+ gt_info = []
669
+ for g in gt_label:
670
+ pts = g.get("points", [])
671
+ if not pts or len(pts) < 4:
672
+ continue
673
+ is_ignored = g.get("transcription", "") == "###"
674
+ gt_info.append({"points": pts, "ignore": is_ignored})
675
+
676
+ # Prepare DET
677
+ det_info = [{"points": d.tolist()} for d in det_boxes]
678
+
679
+ # Evaluate
680
+ result = evaluator.evaluate_image(gt_info, det_info)
681
+ total_gt += result["gt_care"]
682
+ total_det += result["det_care"]
683
+ total_matched += result["det_matched"]
684
+
685
+ per_sample.append({
686
+ "image": img_path,
687
+ "gt_care": result["gt_care"],
688
+ "det_care": result["det_care"],
689
+ "det_matched": result["det_matched"],
690
+ "det_boxes": [d.tolist() for d in det_boxes],
691
+ })
692
+
693
+ if verbose:
694
+ miss = result["gt_care"] - result["det_matched"]
695
+ extra = result["det_care"] - result["det_matched"]
696
+ parts = [
697
+ f"gt={result['gt_care']}",
698
+ f"det={result['det_care']}",
699
+ f"match={result['det_matched']}",
700
+ ]
701
+ if miss > 0:
702
+ parts.append(f"MISS={miss}")
703
+ if extra > 0:
704
+ parts.append(f"EXTRA={extra}")
705
+ print(f"[{os.path.basename(img_path)}] " + " ".join(parts))
706
+
707
+ metrics = evaluator.combine_results(
708
+ [r for r in per_sample if "error" not in r]
709
+ )
710
+
711
+ return {
712
+ "precision": round(metrics["precision"], 6),
713
+ "recall": round(metrics["recall"], 6),
714
+ "hmean": round(metrics["hmean"], 6),
715
+ "total_images": len(samples),
716
+ "total_gt": total_gt,
717
+ "total_det": total_det,
718
+ "det_matched": total_matched,
719
+ "per_sample": per_sample,
720
+ }
721
+
722
+
723
+ # ============================================================================
724
+ # 7. CLI
725
+ # ============================================================================
726
+
727
+ def main():
728
+ parser = argparse.ArgumentParser(
729
+ description="PP-OCRv6 Detection ONNX – inference & evaluation"
730
+ )
731
+ # Model
732
+ parser.add_argument("--det_onnx", type=str,
733
+ default="onnx/det_inference_static_sim.onnx",
734
+ help="Path to detection ONNX model")
735
+ parser.add_argument("--limit_side_len", type=int, default=960)
736
+ parser.add_argument("--det_db_thresh", type=float, default=0.2)
737
+ parser.add_argument("--det_db_box_thresh", type=float, default=0.45)
738
+ parser.add_argument("--det_db_unclip_ratio", type=float, default=1.4)
739
+ parser.add_argument("--use_gpu", action="store_true", help="Enable GPU inference")
740
+ parser.add_argument("--resize_mode", type=str, default="letterbox",
741
+ choices=["letterbox", "stretch"],
742
+ help="Resize strategy for fixed-size ONNX: letterbox (keep ratio+pad) or stretch (direct resize)")
743
+
744
+ # Single image mode
745
+ parser.add_argument("--image", type=str, default=None, help="Single image path")
746
+
747
+ # Evaluation mode
748
+ parser.add_argument("--label_file", type=str,
749
+ default='dataset/ocr_det_dataset_examples/val.txt',
750
+ help="Label file (image_path<TAB>json_label per line)")
751
+ parser.add_argument("--dataset_root", type=str,
752
+ default="dataset/ocr_det_dataset_examples",
753
+ help="Prefix directory for image paths in label file")
754
+
755
+ # Common
756
+ parser.add_argument("--visualize", action="store_true", help="Draw boxes on image")
757
+ parser.add_argument("--output", type=str, default=None,
758
+ help="Save visualized image (implies --visualize)")
759
+ parser.add_argument("--verbose", action="store_true", help="Print per-image metrics")
760
+ parser.add_argument("--output_json", type=str, default=None,
761
+ help="Save results to JSON file")
762
+
763
+ args = parser.parse_args()
764
+
765
+ det = PPOCRv6DetOnnx(
766
+ det_onnx=args.det_onnx,
767
+ det_limit_side_len=args.limit_side_len,
768
+ det_db_thresh=args.det_db_thresh,
769
+ det_db_box_thresh=args.det_db_box_thresh,
770
+ det_db_unclip_ratio=args.det_db_unclip_ratio,
771
+ use_gpu=args.use_gpu,
772
+ resize_mode=args.resize_mode,
773
+ )
774
+
775
+ # --- Single image mode ---
776
+ if args.image and not args.label_file:
777
+ img = cv2.imread(args.image)
778
+ if img is None:
779
+ raise FileNotFoundError(f"Cannot read: {args.image}")
780
+ boxes = det(img)
781
+ print(f"Detected {len(boxes)} text boxes:")
782
+ for i, box in enumerate(boxes):
783
+ print(f" [{i}] {box.tolist()}")
784
+
785
+ do_viz = args.visualize or args.output
786
+ if do_viz:
787
+ vis = draw_det_result(img, boxes)
788
+ out_path = args.output or "det_result.jpg"
789
+ cv2.imwrite(out_path, vis)
790
+ print(f"Visualization saved to: {out_path}")
791
+
792
+ if args.output_json:
793
+ with open(args.output_json, "w") as f:
794
+ json.dump(
795
+ {"image": args.image, "boxes": [b.tolist() for b in boxes]},
796
+ f, indent=2,
797
+ )
798
+ print(f"Results saved to: {args.output_json}")
799
+ return
800
+
801
+ # --- Evaluation mode ---
802
+ if args.label_file:
803
+ metrics = evaluate(
804
+ det, args.label_file,
805
+ dataset_root=args.dataset_root,
806
+ verbose=args.verbose,
807
+ )
808
+
809
+ print()
810
+ print("=" * 60)
811
+ print("Evaluation Results")
812
+ print("=" * 60)
813
+ print(f" Images: {metrics['total_images']}")
814
+ print(f" GT boxes: {metrics['total_gt']}")
815
+ print(f" DET boxes: {metrics['total_det']}")
816
+ print(f" Matched: {metrics['det_matched']}")
817
+ print(f" Precision: {metrics['precision']:.4f} ({metrics['precision']*100:.2f}%)")
818
+ print(f" Recall: {metrics['recall']:.4f} ({metrics['recall']*100:.2f}%)")
819
+ print(f" Hmean (F1): {metrics['hmean']:.4f}")
820
+ print("=" * 60)
821
+
822
+ # Visualization for eval mode
823
+ do_viz = args.visualize or args.output
824
+ if do_viz:
825
+ out_dir = args.output if args.output else "det_eval_vis"
826
+ os.makedirs(out_dir, exist_ok=True)
827
+ for i, smp in enumerate(metrics["per_sample"]):
828
+ img = cv2.imread(smp["image"])
829
+ if img is None:
830
+ continue
831
+ # Load GT boxes with matching info
832
+ with open(args.label_file, "r") as f:
833
+ lines = f.readlines()
834
+ gt_label = []
835
+ for line in lines:
836
+ line = line.strip()
837
+ if not line:
838
+ continue
839
+ parts = line.split("\t")
840
+ if len(parts) < 2:
841
+ continue
842
+ if os.path.join(args.dataset_root, parts[0].strip()) == smp["image"]:
843
+ gt_label = json.loads(parts[1])
844
+ break
845
+
846
+ det_boxes = np.array(smp.get("det_boxes", []))
847
+ # Simple matching for visualization (re-run evaluate_image)
848
+ gt_info = []
849
+ for g in gt_label:
850
+ pts = g.get("points", [])
851
+ if not pts or len(pts) < 4:
852
+ continue
853
+ gt_info.append({
854
+ "points": pts,
855
+ "ignore": g.get("transcription", "") == "###",
856
+ })
857
+ det_info = [{"points": d} for d in det_boxes.tolist()]
858
+
859
+ # Compute matches for coloring
860
+ matched_pairs = _compute_matched_pairs(gt_info, det_info)
861
+ vis = draw_det_result(img, det_boxes, gt_info, matched_pairs)
862
+ fname = os.path.basename(smp["image"])
863
+ cv2.imwrite(os.path.join(out_dir, fname), vis)
864
+
865
+ print(f"Visualization saved to: {out_dir}/")
866
+
867
+ if args.output_json:
868
+ out = {k: v for k, v in metrics.items() if k != "per_sample"}
869
+ out["per_sample"] = metrics["per_sample"]
870
+ with open(args.output_json, "w") as f:
871
+ json.dump(out, f, indent=2)
872
+ print(f"Results saved to: {args.output_json}")
873
+ return
874
+
875
+ parser.error("Either --image or --label_file must be provided.")
876
+
877
+
878
+ def _compute_matched_pairs(gt_info, det_info, iou_thr=0.5):
879
+ """Compute matched GT-det pairs for visualization coloring."""
880
+ pairs = []
881
+ gt_pols = [g["points"] for g in gt_info if not g.get("ignore")]
882
+ det_pols = [d["points"] for d in det_info]
883
+ if not gt_pols or not det_pols:
884
+ return pairs
885
+
886
+ iou_mat = np.empty([len(gt_pols), len(det_pols)])
887
+ for g_i, g_pts in enumerate(gt_pols):
888
+ for d_i, d_pts in enumerate(det_pols):
889
+ try:
890
+ int_area = _get_intersection(d_pts, g_pts)
891
+ union_area = _get_union(d_pts, g_pts)
892
+ iou_mat[g_i, d_i] = int_area / union_area if union_area > 0 else 0
893
+ except Exception:
894
+ iou_mat[g_i, d_i] = 0
895
+
896
+ gt_used = set()
897
+ det_used = set()
898
+ # Greedy matching by descending IoU
899
+ flat = []
900
+ for g_i in range(len(gt_pols)):
901
+ for d_i in range(len(det_pols)):
902
+ flat.append((iou_mat[g_i, d_i], g_i, d_i))
903
+ flat.sort(key=lambda x: x[0], reverse=True)
904
+ for iou, g_i, d_i in flat:
905
+ if iou > iou_thr and g_i not in gt_used and d_i not in det_used:
906
+ pairs.append({"gt": g_i, "det": d_i})
907
+ gt_used.add(g_i)
908
+ det_used.add(d_i)
909
+ return pairs
910
+
911
+
912
+ if __name__ == "__main__":
913
+ main()
run_rec_ax.py ADDED
@@ -0,0 +1,377 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ PP-OCRv6 Recognition AX Inference (one-by-one, batch=1 only)
4
+
5
+ Dependencies:
6
+ numpy, opencv-python, axengine, pyyaml, tqdm
7
+
8
+ Usage:
9
+ # Single image inference
10
+ python ppocrv6_rec_ax-one.py --rec_onnx rec.axmodel --char_dict inference.yml --image crop.png
11
+
12
+ # Batch evaluation
13
+ python ppocrv6_rec_ax-one.py --rec_onnx rec.axmodel --char_dict inference.yml \\
14
+ --label_file labels.txt --dataset_root ./crops/ \\
15
+ --output_json result.json --verbose
16
+ """
17
+
18
+ import argparse
19
+ import math
20
+ import os
21
+ from typing import List, Optional, Tuple, Union
22
+
23
+ import cv2
24
+ import numpy as np
25
+ import axengine as ax
26
+ import yaml
27
+ from tqdm import tqdm
28
+
29
+
30
+ # ============================================================================
31
+ # 1. Utilities
32
+ # ============================================================================
33
+
34
+ def _edit_distance(pred: str, target: str) -> Tuple[int, float]:
35
+ m, n = len(pred), len(target)
36
+ if m == 0:
37
+ return n, 1.0
38
+ if n == 0:
39
+ return m, 1.0
40
+ dp = list(range(n + 1))
41
+ for i in range(1, m + 1):
42
+ prev = dp[0]
43
+ dp[0] = i
44
+ for j in range(1, n + 1):
45
+ temp = dp[j]
46
+ if pred[i - 1] == target[j - 1]:
47
+ dp[j] = prev
48
+ else:
49
+ dp[j] = 1 + min(prev, dp[j], dp[j - 1])
50
+ prev = temp
51
+ distance = dp[n]
52
+ normalized = distance / max(m, n)
53
+ return distance, normalized
54
+
55
+
56
+ def _load_char_dict(source: Union[str, List[str]]) -> List[str]:
57
+ if isinstance(source, list):
58
+ return list(source)
59
+ ext = os.path.splitext(source)[1].lower()
60
+ if ext in (".yml", ".yaml"):
61
+ with open(source, "r", encoding="utf-8") as f:
62
+ cfg = yaml.safe_load(f)
63
+ dic = cfg.get("PostProcess", {}).get("character_dict", [])
64
+ if not dic:
65
+ raise ValueError(f"No PostProcess.character_dict found in {source}")
66
+ return dic
67
+ elif ext == ".txt":
68
+ with open(source, "r", encoding="utf-8") as f:
69
+ return [line.strip("\n\r") for line in f.readlines()]
70
+ else:
71
+ raise ValueError(
72
+ f"Unsupported char_dict source: {source}. Use .yml, .txt, or list."
73
+ )
74
+
75
+
76
+ def _resize_norm_img(
77
+ img: np.ndarray,
78
+ image_shape: Tuple[int, int, int] = (3, 48, 320),
79
+ ) -> np.ndarray:
80
+ imgC, imgH, imgW = image_shape
81
+ h, w = img.shape[:2]
82
+ max_wh_ratio = imgW * 1.0 / imgH
83
+ ratio = w * 1.0 / h
84
+ max_wh_ratio = max(max_wh_ratio, ratio)
85
+ max_wh_ratio = min(max_wh_ratio, imgW / imgH)
86
+ target_w = int(imgH * max_wh_ratio)
87
+
88
+ if math.ceil(imgH * ratio) > target_w:
89
+ resized_w = target_w
90
+ else:
91
+ resized_w = int(math.ceil(imgH * ratio))
92
+
93
+ resized = cv2.resize(img, (resized_w, imgH))
94
+ resized = resized.astype("float32")
95
+ resized = resized.transpose((2, 0, 1))
96
+ # resized /= 255.0
97
+ # resized -= 0.5
98
+ # resized /= 0.5
99
+ padded = np.zeros((imgC, imgH, target_w), dtype=np.float32)
100
+ padded[:, :, 0:resized_w] = resized
101
+ return padded
102
+
103
+
104
+ # ============================================================================
105
+ # 2. CTC Decoder
106
+ # ============================================================================
107
+
108
+ class _CTCLabelDecode:
109
+ """CTC greedy decoder for recognition output."""
110
+
111
+ def __init__(self, character_list: List[str], use_space_char: bool = True):
112
+ self.character_str = list(character_list)
113
+ if use_space_char:
114
+ self.character_str.append(" ")
115
+ dict_character = ["blank"] + self.character_str
116
+ self.character = dict_character
117
+ self.dict = {char: i for i, char in enumerate(dict_character)}
118
+
119
+ def decode(
120
+ self,
121
+ text_index: np.ndarray,
122
+ text_prob: Optional[np.ndarray] = None,
123
+ is_remove_duplicate: bool = True,
124
+ ) -> List[Tuple[str, float]]:
125
+ result_list = []
126
+ batch_size = len(text_index)
127
+ for batch_idx in range(batch_size):
128
+ selection = np.ones(len(text_index[batch_idx]), dtype=bool)
129
+ if is_remove_duplicate:
130
+ selection[1:] = text_index[batch_idx][1:] != text_index[batch_idx][:-1]
131
+ selection &= text_index[batch_idx] != 0
132
+ char_list = [
133
+ self.character[int(tid)]
134
+ for tid in text_index[batch_idx][selection]
135
+ ]
136
+ if text_prob is not None:
137
+ conf_list = text_prob[batch_idx][selection]
138
+ else:
139
+ conf_list = np.ones(len(selection), dtype=np.float32)
140
+ if len(conf_list) == 0:
141
+ conf_list = np.array([0.0], dtype=np.float32)
142
+ text = "".join(char_list)
143
+ result_list.append((text, float(np.mean(conf_list))))
144
+ return result_list
145
+
146
+ def __call__(self, preds: np.ndarray) -> List[Tuple[str, float]]:
147
+ preds_idx = preds.argmax(axis=2)
148
+ preds_prob = preds.max(axis=2)
149
+ return self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
150
+
151
+
152
+ # ============================================================================
153
+ # 3. Recognition Engine (one-by-one)
154
+ # ============================================================================
155
+
156
+ class PPOCRv6RecOne:
157
+ """Recognition engine for batch=1 AX model.
158
+
159
+ Flow: load image → preprocess single → inference → decode.
160
+ """
161
+
162
+ def __init__(
163
+ self,
164
+ rec_model: str,
165
+ char_dict: Union[str, List[str]],
166
+ rec_image_shape: Tuple[int, int, int] = (3, 48, 320),
167
+ ):
168
+ self.rec_image_shape = rec_image_shape
169
+ self.session = ax.InferenceSession(rec_model)
170
+ self.input_name = self.session.get_inputs()[0].name
171
+ char_list = _load_char_dict(char_dict)
172
+ self._decoder = _CTCLabelDecode(char_list, use_space_char=True)
173
+
174
+ def _preprocess_one(self, img: np.ndarray) -> np.ndarray:
175
+ tensor = _resize_norm_img(img, self.rec_image_shape)
176
+ return np.expand_dims(tensor, axis=0).astype(np.float32)
177
+
178
+ def infer_one(self, img: np.ndarray) -> Tuple[str, float]:
179
+ inp = self._preprocess_one(img)
180
+ out = self.session.run(None, {self.input_name: inp})[0]
181
+ results = self._decoder(out)
182
+ return results[0] if results else ("", 0.0)
183
+
184
+ def __call__(
185
+ self, images: List[np.ndarray]
186
+ ) -> List[Tuple[str, float]]:
187
+ results = []
188
+ for img in tqdm(images, desc="Inference", unit="sample"):
189
+ results.append(self.infer_one(img))
190
+ return results
191
+
192
+ def predict_image(self, path: str) -> Tuple[str, float]:
193
+ im = cv2.imread(path)
194
+ if im is None:
195
+ raise FileNotFoundError(f"Cannot read: {path}")
196
+ return self.infer_one(im)
197
+
198
+
199
+ # ============================================================================
200
+ # 4. Evaluation (one-by-one)
201
+ # ============================================================================
202
+
203
+ def evaluate(
204
+ ocr: PPOCRv6RecOne,
205
+ label_file: str,
206
+ dataset_root: str = "",
207
+ ignore_space: bool = True,
208
+ verbose: bool = False,
209
+ ) -> dict:
210
+ images = []
211
+ targets = []
212
+
213
+ with open(label_file, "r", encoding="utf-8") as f:
214
+ for line in f:
215
+ line = line.strip()
216
+ if not line:
217
+ continue
218
+ parts = line.split("\t")
219
+ if len(parts) < 2:
220
+ continue
221
+ img_path = os.path.join(dataset_root, parts[0].strip())
222
+ images.append(img_path)
223
+ targets.append(parts[1].strip())
224
+
225
+ total = len(images)
226
+ if total == 0:
227
+ print("[WARN] No samples found in label file.")
228
+ return {"acc": 0.0, "norm_edit_dis": 0.0, "total": 0, "correct": 0, "per_sample": []}
229
+
230
+ correct = 0
231
+ total_edit_dis = 0.0
232
+ per_sample = []
233
+
234
+ for i in tqdm(range(total), desc="Evaluating", unit="sample"):
235
+ p = images[i]
236
+ gt = targets[i]
237
+ im = cv2.imread(p)
238
+ if im is None:
239
+ print(f"[WARN] Cannot read {p}, skipping.")
240
+ pred, conf = ("", 0.0)
241
+ else:
242
+ pred, conf = ocr.infer_one(im)
243
+
244
+ pred_clean = pred.replace(" ", "") if ignore_space else pred
245
+ gt_clean = gt.replace(" ", "") if ignore_space else gt
246
+ dist, norm_dist = _edit_distance(pred_clean, gt_clean)
247
+ total_edit_dis += norm_dist
248
+ is_correct = pred_clean == gt_clean
249
+ if is_correct:
250
+ correct += 1
251
+
252
+ sample = {
253
+ "image": p,
254
+ "pred": pred,
255
+ "gt": gt,
256
+ "confidence": round(conf, 4),
257
+ "correct": is_correct,
258
+ "edit_distance": int(dist),
259
+ "norm_edit_dis": round(norm_dist, 4),
260
+ }
261
+ per_sample.append(sample)
262
+
263
+ if verbose:
264
+ status = " OK" if is_correct else "MIS"
265
+ print(
266
+ f"[{status}] pred={pred!r:<30} gt={gt!r:<30} "
267
+ f"conf={conf:.4f} edit={int(dist)} ndis={norm_dist:.4f}"
268
+ )
269
+
270
+ acc = correct / total
271
+ norm_edit_dis = 1.0 - total_edit_dis / total
272
+
273
+ return {
274
+ "acc": round(acc, 6),
275
+ "norm_edit_dis": round(norm_edit_dis, 6),
276
+ "total": total,
277
+ "correct": correct,
278
+ "per_sample": per_sample,
279
+ }
280
+
281
+
282
+ # ============================================================================
283
+ # 5. CLI
284
+ # ============================================================================
285
+
286
+ def main():
287
+ parser = argparse.ArgumentParser(
288
+ description="PP-OCRv6 Recognition AX (one-by-one) – inference & evaluation"
289
+ )
290
+ parser.add_argument(
291
+ "--rec_onnx", type=str,
292
+ default="axmodel/ax650/rec_npu3.axmodel",
293
+ help="Path to recognition AX model",
294
+ )
295
+ parser.add_argument(
296
+ "--char_dict", type=str,
297
+ default="onnx/rec_inference.yml",
298
+ help="Character dictionary: .yml (PostProcess.character_dict), .txt, or comma-list",
299
+ )
300
+ parser.add_argument("--rec_image_shape", type=str, default="3,48,320",
301
+ help="Recognition input shape C,H,W (comma separated)")
302
+
303
+ parser.add_argument("--image", type=str, default=None, help="Single crop image path")
304
+
305
+ parser.add_argument("--label_file", type=str,
306
+ default='dataset/ocr_rec_dataset_examples/val.txt',
307
+ help="Label file (image_path<TAB>gt_text per line)")
308
+ parser.add_argument("--dataset_root", type=str,
309
+ default="dataset/ocr_rec_dataset_examples",
310
+ help="Prefix directory for image paths in label file")
311
+
312
+ parser.add_argument("--ignore_space", action="store_true", default=True,
313
+ help="Ignore spaces when comparing (default: True)")
314
+ parser.add_argument("--verbose", action="store_true", help="Print per-sample results")
315
+ parser.add_argument("--output_json", type=str, default=None,
316
+ help="Save results to JSON file")
317
+
318
+ args = parser.parse_args()
319
+
320
+ char_dict_src = args.char_dict
321
+ if char_dict_src.startswith("[") or ("," in char_dict_src and not os.path.exists(char_dict_src)):
322
+ char_dict = [c.strip() for c in char_dict_src.split(",") if c.strip()]
323
+ else:
324
+ char_dict = char_dict_src
325
+
326
+ image_shape = tuple(int(v) for v in args.rec_image_shape.split(","))
327
+ if len(image_shape) != 3:
328
+ raise ValueError("--rec_image_shape requires 3 comma-separated integers")
329
+
330
+ ocr = PPOCRv6RecOne(
331
+ rec_model=args.rec_onnx,
332
+ char_dict=char_dict,
333
+ rec_image_shape=image_shape,
334
+ )
335
+
336
+ if args.image:
337
+ text, conf = ocr.predict_image(args.image)
338
+ print(f"text={text!r} confidence={conf:.4f}")
339
+ if args.output_json:
340
+ import json
341
+ with open(args.output_json, "w", encoding="utf-8") as f:
342
+ json.dump({"text": text, "confidence": conf}, f, ensure_ascii=False, indent=2)
343
+ return
344
+
345
+ if args.label_file:
346
+ metrics = evaluate(
347
+ ocr,
348
+ args.label_file,
349
+ dataset_root=args.dataset_root,
350
+ ignore_space=args.ignore_space,
351
+ verbose=args.verbose,
352
+ )
353
+
354
+ print()
355
+ print("=" * 60)
356
+ print("Evaluation Results")
357
+ print("=" * 60)
358
+ print(f" Total samples: {metrics['total']}")
359
+ print(f" Correct (exact match): {metrics['correct']}")
360
+ print(f" Accuracy: {metrics['acc']:.4f} ({metrics['acc']*100:.2f}%)")
361
+ print(f" Norm Edit Distance: {metrics['norm_edit_dis']:.4f}")
362
+ print("=" * 60)
363
+
364
+ if args.output_json:
365
+ import json
366
+ out = {k: v for k, v in metrics.items() if k != "per_sample"}
367
+ out["per_sample"] = metrics["per_sample"]
368
+ with open(args.output_json, "w", encoding="utf-8") as f:
369
+ json.dump(out, f, ensure_ascii=False, indent=2)
370
+ print(f"\nResults saved to: {args.output_json}")
371
+ return
372
+
373
+ parser.error("Either --image or --label_file must be provided.")
374
+
375
+
376
+ if __name__ == "__main__":
377
+ main()
run_rec_onnx.py ADDED
@@ -0,0 +1,494 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ PP-OCRv6 Recognition ONNX Inference & Evaluation (standalone, zero Paddle dependency)
4
+
5
+ Dependencies:
6
+ numpy, opencv-python, onnxruntime, pyyaml
7
+
8
+ Usage:
9
+ # Single image inference
10
+ python ppocrv6_rec_onnx.py --rec_onnx rec.onnx --char_dict inference.yml --image crop.png
11
+
12
+ # Batch evaluation
13
+ python ppocrv6_rec_onnx.py --rec_onnx rec.onnx --char_dict inference.yml \\
14
+ --label_file labels.txt --dataset_root ./crops/ \\
15
+ --batch_size 8 --output_json result.json --verbose
16
+ """
17
+
18
+ import argparse
19
+ import math
20
+ import os
21
+ from typing import List, Optional, Tuple, Union
22
+
23
+ import cv2
24
+ import numpy as np
25
+ import onnxruntime as ort
26
+ import yaml
27
+
28
+
29
+ # ============================================================================
30
+ # 1. Utilities
31
+ # ============================================================================
32
+
33
+ def _edit_distance(pred: str, target: str) -> Tuple[int, float]:
34
+ """Compute Levenshtein edit distance (pure Python, no extra deps).
35
+
36
+ Returns:
37
+ (distance, normalized_distance) where normalized ∈ [0, 1].
38
+ """
39
+ m, n = len(pred), len(target)
40
+ if m == 0:
41
+ return n, 1.0
42
+ if n == 0:
43
+ return m, 1.0
44
+
45
+ dp = list(range(n + 1))
46
+ for i in range(1, m + 1):
47
+ prev = dp[0]
48
+ dp[0] = i
49
+ for j in range(1, n + 1):
50
+ temp = dp[j]
51
+ if pred[i - 1] == target[j - 1]:
52
+ dp[j] = prev
53
+ else:
54
+ dp[j] = 1 + min(prev, dp[j], dp[j - 1])
55
+ prev = temp
56
+
57
+ distance = dp[n]
58
+ normalized = distance / max(m, n)
59
+ return distance, normalized
60
+
61
+
62
+ def _load_char_dict(source: Union[str, List[str]]) -> List[str]:
63
+ """Load character dictionary from .yml, .txt, or list."""
64
+ if isinstance(source, list):
65
+ return list(source)
66
+ ext = os.path.splitext(source)[1].lower()
67
+ if ext in (".yml", ".yaml"):
68
+ with open(source, "r", encoding="utf-8") as f:
69
+ cfg = yaml.safe_load(f)
70
+ dic = cfg.get("PostProcess", {}).get("character_dict", [])
71
+ if not dic:
72
+ raise ValueError(f"No PostProcess.character_dict found in {source}")
73
+ return dic
74
+ elif ext == ".txt":
75
+ with open(source, "r", encoding="utf-8") as f:
76
+ return [line.strip("\n\r") for line in f.readlines()]
77
+ else:
78
+ raise ValueError(
79
+ f"Unsupported char_dict source: {source}. Use .yml, .txt, or list."
80
+ )
81
+
82
+
83
+ def _resize_norm_img(
84
+ img: np.ndarray,
85
+ image_shape: Tuple[int, int, int] = (3, 48, 320),
86
+ max_wh_ratio: Optional[float] = None,
87
+ ) -> np.ndarray:
88
+ """Resize and normalize a cropped text image for recognition.
89
+
90
+ Args:
91
+ img: BGR crop image (H, W, 3).
92
+ image_shape: (C, H, W) target shape.
93
+ max_wh_ratio: precomputed max width/height ratio for batch. If None, derived from img.
94
+ """
95
+ imgC, imgH, imgW = image_shape
96
+ if max_wh_ratio is None:
97
+ max_wh_ratio = imgW * 1.0 / imgH
98
+ h, w = img.shape[:2]
99
+ ratio = w * 1.0 / h
100
+ max_wh_ratio = max(max_wh_ratio, ratio)
101
+
102
+ max_wh_ratio = min(max_wh_ratio, imgW / imgH)
103
+ target_w = int(imgH * max_wh_ratio)
104
+ h, w = img.shape[:2]
105
+ ratio = w * 1.0 / h
106
+ if math.ceil(imgH * ratio) > target_w:
107
+ resized_w = target_w
108
+ else:
109
+ resized_w = int(math.ceil(imgH * ratio))
110
+
111
+ resized = cv2.resize(img, (resized_w, imgH))
112
+ resized = resized.astype("float32")
113
+ resized = resized.transpose((2, 0, 1)) / 255.0
114
+ resized -= 0.5
115
+ resized /= 0.5
116
+ padded = np.zeros((imgC, imgH, target_w), dtype=np.float32)
117
+ padded[:, :, 0:resized_w] = resized
118
+ return padded
119
+
120
+
121
+ # ============================================================================
122
+ # 2. CTC Decoder
123
+ # ============================================================================
124
+
125
+ class _CTCLabelDecode:
126
+ """CTC greedy decoder for recognition output."""
127
+
128
+ def __init__(self, character_list: List[str], use_space_char: bool = True):
129
+ self.character_str = list(character_list)
130
+ if use_space_char:
131
+ self.character_str.append(" ")
132
+ dict_character = ["blank"] + self.character_str
133
+ self.character = dict_character
134
+ self.dict = {char: i for i, char in enumerate(dict_character)}
135
+
136
+ def decode(
137
+ self,
138
+ text_index: np.ndarray,
139
+ text_prob: Optional[np.ndarray] = None,
140
+ is_remove_duplicate: bool = True,
141
+ ) -> List[Tuple[str, float]]:
142
+ result_list = []
143
+ batch_size = len(text_index)
144
+ for batch_idx in range(batch_size):
145
+ selection = np.ones(len(text_index[batch_idx]), dtype=bool)
146
+ if is_remove_duplicate:
147
+ selection[1:] = text_index[batch_idx][1:] != text_index[batch_idx][:-1]
148
+ selection &= text_index[batch_idx] != 0 # ignore blank
149
+ char_list = [
150
+ self.character[int(tid)]
151
+ for tid in text_index[batch_idx][selection]
152
+ ]
153
+ if text_prob is not None:
154
+ conf_list = text_prob[batch_idx][selection]
155
+ else:
156
+ conf_list = np.ones(len(selection), dtype=np.float32)
157
+ if len(conf_list) == 0:
158
+ conf_list = np.array([0.0], dtype=np.float32)
159
+ text = "".join(char_list)
160
+ result_list.append((text, float(np.mean(conf_list))))
161
+ return result_list
162
+
163
+ def __call__(self, preds: np.ndarray) -> List[Tuple[str, float]]:
164
+ preds_idx = preds.argmax(axis=2)
165
+ preds_prob = preds.max(axis=2)
166
+ return self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
167
+
168
+
169
+ # ============================================================================
170
+ # 3. Recognition Engine
171
+ # ============================================================================
172
+
173
+ class PPOCRv6RecOnnx:
174
+
175
+ def __init__(
176
+ self,
177
+ rec_onnx: str,
178
+ char_dict: Union[str, List[str]],
179
+ rec_image_shape: Tuple[int, int, int] = (3, 48, 320),
180
+ rec_batch_num: int = 6,
181
+ use_gpu: bool = False,
182
+ onnx_providers: Optional[List[str]] = None,
183
+ ):
184
+ self.rec_image_shape = rec_image_shape
185
+ self.rec_batch_num = rec_batch_num
186
+
187
+ # ONNX session
188
+ if onnx_providers is None:
189
+ onnx_providers = (
190
+ ["CUDAExecutionProvider", "CPUExecutionProvider"]
191
+ if use_gpu
192
+ else ["CPUExecutionProvider"]
193
+ )
194
+
195
+ sess_options = ort.SessionOptions()
196
+ sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
197
+ self.session = ort.InferenceSession(
198
+ rec_onnx, sess_options=sess_options, providers=onnx_providers
199
+ )
200
+ self.input_name = self.session.get_inputs()[0].name
201
+
202
+ # CTC decoder
203
+ char_list = _load_char_dict(char_dict)
204
+ self._decoder = _CTCLabelDecode(char_list, use_space_char=True)
205
+
206
+ # ---- pre / post ----
207
+
208
+ def _preprocess(
209
+ self, img_list: List[np.ndarray]
210
+ ) -> List[np.ndarray]:
211
+ """Convert a list of crops into batch tensors (grouped by self.rec_batch_num)."""
212
+ num = len(img_list)
213
+ width_list = [im.shape[1] / float(im.shape[0]) for im in img_list]
214
+ indices = np.argsort(np.array(width_list))
215
+
216
+ batches = []
217
+ index_maps = [] # each element: list of original indices in this batch
218
+
219
+ for beg in range(0, num, self.rec_batch_num):
220
+ end = min(num, beg + self.rec_batch_num)
221
+ imgC, imgH, imgW = self.rec_image_shape
222
+ max_wh_ratio = imgW / imgH
223
+ for ino in range(beg, end):
224
+ orig_idx = indices[ino]
225
+ h, w = img_list[orig_idx].shape[:2]
226
+ max_wh_ratio = max(max_wh_ratio, w / h)
227
+
228
+ norm_list = []
229
+ idx_list = []
230
+ for ino in range(beg, end):
231
+ orig_idx = indices[ino]
232
+ norm = _resize_norm_img(
233
+ img_list[orig_idx],
234
+ self.rec_image_shape,
235
+ max_wh_ratio=max_wh_ratio,
236
+ )
237
+ norm_list.append(np.expand_dims(norm, axis=0))
238
+ idx_list.append(orig_idx)
239
+
240
+ if norm_list:
241
+ batches.append(np.concatenate(norm_list, axis=0).astype(np.float32))
242
+ index_maps.append(idx_list)
243
+ return batches, index_maps
244
+
245
+ def _postprocess(
246
+ self,
247
+ batch_outputs: List[np.ndarray],
248
+ index_maps: List[List[int]],
249
+ total_num: int,
250
+ ) -> List[Tuple[str, float]]:
251
+ results = [("", 0.0)] * total_num
252
+ # Decode each batch separately — different batches may have different T.
253
+ for batch_preds, idx_list in zip(batch_outputs, index_maps):
254
+ texts = self._decoder(batch_preds)
255
+ for i, orig_idx in enumerate(idx_list):
256
+ results[orig_idx] = texts[i]
257
+ return results
258
+
259
+ # ---- public API ----
260
+
261
+ def __call__(
262
+ self, img: Union[np.ndarray, List[np.ndarray]]
263
+ ) -> List[Tuple[str, float]]:
264
+ if isinstance(img, np.ndarray):
265
+ img = [img]
266
+ if not img:
267
+ return []
268
+ batches, index_maps = self._preprocess(img)
269
+ outputs = []
270
+ for batch in batches:
271
+ out = self.session.run(None, {self.input_name: batch})
272
+ outputs.append(out[0])
273
+ return self._postprocess(outputs, index_maps, len(img))
274
+
275
+ def predict_image(self, path: str) -> Tuple[str, float]:
276
+ im = cv2.imread(path)
277
+ if im is None:
278
+ raise FileNotFoundError(f"Cannot read: {path}")
279
+ return self.__call__(im)[0]
280
+
281
+
282
+ # ============================================================================
283
+ # 4. Evaluation
284
+ # ============================================================================
285
+
286
+ def evaluate(
287
+ ocr: PPOCRv6RecOnnx,
288
+ label_file: str,
289
+ dataset_root: str = "",
290
+ ignore_space: bool = True,
291
+ verbose: bool = False,
292
+ ) -> dict:
293
+ """Evaluate recognition accuracy against a ground-truth label file.
294
+
295
+ Label file format (one per line, tab-separated)::
296
+
297
+ rel/path/to/crop.png<TAB>ground truth text
298
+
299
+ The full image path is ``os.path.join(dataset_root, rel_path)``.
300
+
301
+ Args:
302
+ ocr: PPOCRv6RecOnnx instance.
303
+ label_file: path to tab-separated label file.
304
+ dataset_root: prefix directory for image paths in label file.
305
+ ignore_space: strip spaces before comparing.
306
+ verbose: print per-sample prediction details.
307
+
308
+ Returns:
309
+ dict with keys: ``acc``, ``norm_edit_dis``, ``total``, ``correct``,
310
+ ``per_sample`` (list of per-sample details).
311
+ """
312
+ images = []
313
+ targets = []
314
+
315
+ with open(label_file, "r", encoding="utf-8") as f:
316
+ for line in f:
317
+ line = line.strip()
318
+ if not line:
319
+ continue
320
+ parts = line.split("\t")
321
+ if len(parts) < 2:
322
+ continue
323
+ img_path = os.path.join(dataset_root, parts[0].strip())
324
+ images.append(img_path)
325
+ targets.append(parts[1].strip())
326
+
327
+ total = len(images)
328
+ if total == 0:
329
+ print("[WARN] No samples found in label file.")
330
+ return {"acc": 0.0, "norm_edit_dis": 0.0, "total": 0, "correct": 0, "per_sample": []}
331
+
332
+ # Load all images
333
+ imgs = []
334
+ for p in images:
335
+ im = cv2.imread(p)
336
+ if im is None:
337
+ print(f"[WARN] Cannot read {p}, skipping.")
338
+ imgs.append(np.zeros((32, 100, 3), dtype=np.uint8))
339
+ else:
340
+ imgs.append(im)
341
+
342
+ # Batch inference
343
+ rec_results = ocr(imgs)
344
+
345
+ correct = 0
346
+ total_edit_dis = 0.0
347
+ per_sample = []
348
+
349
+ for i, ((pred, conf), gt) in enumerate(zip(rec_results, targets)):
350
+ pred_clean = pred.replace(" ", "") if ignore_space else pred
351
+ gt_clean = gt.replace(" ", "") if ignore_space else gt
352
+ dist, norm_dist = _edit_distance(pred_clean, gt_clean)
353
+ total_edit_dis += norm_dist
354
+ is_correct = pred_clean == gt_clean
355
+ if is_correct:
356
+ correct += 1
357
+
358
+ sample = {
359
+ "image": images[i],
360
+ "pred": pred,
361
+ "gt": gt,
362
+ "confidence": round(conf, 4),
363
+ "correct": is_correct,
364
+ "edit_distance": int(dist),
365
+ "norm_edit_dis": round(norm_dist, 4),
366
+ }
367
+ per_sample.append(sample)
368
+
369
+ if verbose:
370
+ status = " OK" if is_correct else "MIS"
371
+ print(
372
+ f"[{status}] pred={pred!r:<30} gt={gt!r:<30} "
373
+ f"conf={conf:.4f} edit={int(dist)} ndis={norm_dist:.4f}"
374
+ )
375
+
376
+ acc = correct / total
377
+ norm_edit_dis = 1.0 - total_edit_dis / total
378
+
379
+ return {
380
+ "acc": round(acc, 6),
381
+ "norm_edit_dis": round(norm_edit_dis, 6),
382
+ "total": total,
383
+ "correct": correct,
384
+ "per_sample": per_sample,
385
+ }
386
+
387
+
388
+ # ============================================================================
389
+ # 5. CLI
390
+ # ============================================================================
391
+
392
+ def main():
393
+ parser = argparse.ArgumentParser(
394
+ description="PP-OCRv6 Recognition ONNX – inference & evaluation"
395
+ )
396
+ # Model
397
+ parser.add_argument(
398
+ "--rec_onnx", type=str,
399
+ default="onnx/rec_inference_static_sim.onnx",
400
+ help="Path to recognition ONNX model",
401
+ )
402
+ parser.add_argument(
403
+ "--char_dict", type=str,
404
+ default="onnx/rec_inference.yml",
405
+ help="Character dictionary: .yml (PostProcess.character_dict), .txt, or comma-list",
406
+ )
407
+ parser.add_argument("--batch_size", type=int, default=1, help="Recognition batch size")
408
+ parser.add_argument("--use_gpu", action="store_true", help="Enable GPU inference")
409
+ parser.add_argument("--rec_image_shape", type=str, default="3,48,320",
410
+ help="Recognition input shape C,H,W (comma separated)")
411
+
412
+ # Single image mode
413
+ parser.add_argument("--image", type=str, default=None, help="Single crop image path")
414
+
415
+ # Evaluation mode
416
+ parser.add_argument("--label_file", type=str,
417
+ default='dataset/ocr_rec_dataset_examples/val.txt',
418
+ help="Label file (image_path<TAB>gt_text per line)")
419
+ parser.add_argument("--dataset_root", type=str,
420
+ default="dataset/ocr_rec_dataset_examples",
421
+ help="Prefix directory for image paths in label file")
422
+
423
+ # Common
424
+ parser.add_argument("--ignore_space", action="store_true", default=True,
425
+ help="Ignore spaces when comparing (default: True)")
426
+ parser.add_argument("--verbose", action="store_true", help="Print per-sample results")
427
+ parser.add_argument("--output_json", type=str, default=None,
428
+ help="Save results to JSON file")
429
+
430
+ args = parser.parse_args()
431
+
432
+ # Build engine
433
+ char_dict_src = args.char_dict
434
+ if char_dict_src.startswith("[") or ("," in char_dict_src and not os.path.exists(char_dict_src)):
435
+ char_dict = [c.strip() for c in char_dict_src.split(",") if c.strip()]
436
+ else:
437
+ char_dict = char_dict_src
438
+
439
+ image_shape = tuple(int(v) for v in args.rec_image_shape.split(","))
440
+ if len(image_shape) != 3:
441
+ raise ValueError("--rec_image_shape requires 3 comma-separated integers")
442
+
443
+ ocr = PPOCRv6RecOnnx(
444
+ rec_onnx=args.rec_onnx,
445
+ char_dict=char_dict,
446
+ rec_image_shape=image_shape,
447
+ rec_batch_num=args.batch_size,
448
+ use_gpu=args.use_gpu,
449
+ )
450
+
451
+ # Single image mode
452
+ if args.image:
453
+ text, conf = ocr.predict_image(args.image)
454
+ print(f"text={text!r} confidence={conf:.4f}")
455
+ if args.output_json:
456
+ import json
457
+ with open(args.output_json, "w", encoding="utf-8") as f:
458
+ json.dump({"text": text, "confidence": conf}, f, ensure_ascii=False, indent=2)
459
+ return
460
+
461
+ # Evaluation mode
462
+ if args.label_file:
463
+ metrics = evaluate(
464
+ ocr,
465
+ args.label_file,
466
+ dataset_root=args.dataset_root,
467
+ ignore_space=args.ignore_space,
468
+ verbose=args.verbose,
469
+ )
470
+
471
+ print()
472
+ print("=" * 60)
473
+ print("Evaluation Results")
474
+ print("=" * 60)
475
+ print(f" Total samples: {metrics['total']}")
476
+ print(f" Correct (exact match): {metrics['correct']}")
477
+ print(f" Accuracy: {metrics['acc']:.4f} ({metrics['acc']*100:.2f}%)")
478
+ print(f" Norm Edit Distance: {metrics['norm_edit_dis']:.4f}")
479
+ print("=" * 60)
480
+
481
+ if args.output_json:
482
+ import json
483
+ out = {k: v for k, v in metrics.items() if k != "per_sample"}
484
+ out["per_sample"] = metrics["per_sample"]
485
+ with open(args.output_json, "w", encoding="utf-8") as f:
486
+ json.dump(out, f, ensure_ascii=False, indent=2)
487
+ print(f"\nResults saved to: {args.output_json}")
488
+ return
489
+
490
+ parser.error("Either --image or --label_file must be provided.")
491
+
492
+
493
+ if __name__ == "__main__":
494
+ main()