Upload folder using huggingface_hub
Browse files- .gitattributes +9 -0
- 11.jpg +0 -0
- README-zh.md +186 -0
- README.md +183 -1
- axmodel/ax650/cls_npu1.axmodel +3 -0
- axmodel/ax650/cls_npu3.axmodel +3 -0
- axmodel/ax650/det_npu1.axmodel +3 -0
- axmodel/ax650/det_npu3.axmodel +3 -0
- axmodel/ax650/rec_npu1.axmodel +3 -0
- axmodel/ax650/rec_npu3.axmodel +3 -0
- cls.json +51 -0
- dataset/download_quant_dataset.sh +2 -0
- dataset/download_val_dataset.sh +2 -0
- det.json +46 -0
- fonts/simfang.ttf +3 -0
- onnx/cls_mobile_sim_static.onnx +3 -0
- onnx/det_inference_static_sim.onnx +3 -0
- onnx/rec_inference.yml +0 -0
- onnx/rec_inference_static_sim.onnx +3 -0
- ppocrv6_ax.py +721 -0
- ppocrv6_onnx.py +721 -0
- rec.json +54 -0
- res-ax.jpg +3 -0
- res-onnx.jpg +3 -0
- run_det_ax.py +894 -0
- run_det_onnx.py +913 -0
- run_rec_ax.py +377 -0
- run_rec_onnx.py +494 -0
.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 |
+

|
| 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 |
+

|
| 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:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+

|
| 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 |
+

|
| 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
|
res-onnx.jpg
ADDED
|
Git LFS Details
|
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()
|