89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""MinerU RDNA 适配补丁 — 一键应用脚本"""
|
|
|
|
import re
|
|
import os
|
|
import sys
|
|
|
|
|
|
def find_infer_dir():
|
|
import mineru.model.utils.tools.infer
|
|
return os.path.dirname(mineru.model.utils.tools.infer.__file__)
|
|
|
|
|
|
def patch_a_predict_rec_imgw(infer_dir):
|
|
"""predict_rec.py: imgW 对齐到 32"""
|
|
f = os.path.join(infer_dir, 'predict_rec.py')
|
|
c = open(f).read()
|
|
old = r'(^ *)(imgW = max\(min\(imgW, self\.limited_max_width\), self\.limited_min_width\)\n)'
|
|
new = r'\1\2\1imgW = math.ceil(imgW / 32) * 32\n'
|
|
c2 = re.sub(old, new, c, flags=re.MULTILINE)
|
|
if c2 == c:
|
|
if 'math.ceil(imgW / 32)' not in c:
|
|
raise RuntimeError('Patch A: cannot find imgW line')
|
|
print('Patch A: already applied')
|
|
else:
|
|
open(f, 'w').write(c2)
|
|
print('Patch A: imgW 32-align inserted')
|
|
|
|
|
|
def patch_b_predict_rec_batch(infer_dir):
|
|
"""predict_rec.py: 批次填充"""
|
|
f = os.path.join(infer_dir, 'predict_rec.py')
|
|
c = open(f).read()
|
|
old = r'(^ *)(norm_img_batch = np\.concatenate\(norm_img_batch\))'
|
|
new = (
|
|
r'\1actual_batch_size = len(norm_img_batch)\n'
|
|
r'\1if actual_batch_size < batch_num:\n'
|
|
r'\1 pad_size = batch_num - actual_batch_size\n'
|
|
r'\1 pad_img = np.zeros_like(norm_img_batch[0])\n'
|
|
r'\1 for _ in range(pad_size):\n'
|
|
r'\1 norm_img_batch.append(pad_img)\n'
|
|
r'\1\2'
|
|
)
|
|
c2 = re.sub(old, new, c, flags=re.MULTILINE)
|
|
if c2 == c:
|
|
if 'actual_batch_size' not in c:
|
|
raise RuntimeError('Patch B: cannot find norm_img_batch concatenation')
|
|
print('Patch B: already applied')
|
|
else:
|
|
open(f, 'w').write(c2)
|
|
print('Patch B: batch padding inserted')
|
|
# 修改 range(len(rec_result)) → range(actual_batch_size)
|
|
c3 = open(f).read()
|
|
c4 = re.sub(
|
|
r'( +)for rno in range\(len\(rec_result\)\):',
|
|
r'\1for rno in range(actual_batch_size):',
|
|
c3
|
|
)
|
|
open(f, 'w').write(c4)
|
|
|
|
|
|
def patch_c_predict_det_contiguous(infer_dir):
|
|
"""predict_det.py: contiguous 检查"""
|
|
f = os.path.join(infer_dir, 'predict_det.py')
|
|
c = open(f).read()
|
|
old = r'(^ *)(inp = inp\.to\(self\.device\)\n)'
|
|
new = r'\1\2\1if not inp.is_contiguous():\n\1 inp = inp.contiguous()\n'
|
|
c2 = re.sub(old, new, c, flags=re.MULTILINE)
|
|
if c2 == c:
|
|
if 'is_contiguous' not in c:
|
|
raise RuntimeError('Patch C: cannot find inp.to(device) line')
|
|
print('Patch C: already applied')
|
|
else:
|
|
open(f, 'w').write(c2)
|
|
print('Patch C: contiguous check inserted')
|
|
|
|
|
|
def main():
|
|
infer_dir = find_infer_dir()
|
|
print(f'MinerU infer dir: {infer_dir}')
|
|
patch_a_predict_rec_imgw(infer_dir)
|
|
patch_b_predict_rec_batch(infer_dir)
|
|
patch_c_predict_det_contiguous(infer_dir)
|
|
print('All MinerU RDNA patches applied.')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|