x

2026-06-03 15:52:29 +08:00
parent 8a36630543
commit b92f26120e
2 changed files with 99 additions and 93 deletions
@@ -0,0 +1,88 @@
+#!/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        imgW = math.ceil(imgW / 32) * 32\n'
+    c2 = re.sub(old, new, c)
+    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 = (
+        '        actual_batch_size = len(norm_img_batch)\n'
+        '        if actual_batch_size < batch_num:\n'
+        '            pad_size = batch_num - actual_batch_size\n'
+        '            pad_img = np.zeros_like(norm_img_batch[0])\n'
+        '            for _ in range(pad_size):\n'
+        '                norm_img_batch.append(pad_img)\n'
+        r'\1'
+    )
+    c2 = re.sub(old, new, c)
+    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\)\):',
+        '                for 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            if not inp.is_contiguous():\n                inp = inp.contiguous()\n'
+    c2 = re.sub(old, new, c)
+    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()