diff --git a/docker/scripts/apply_mineru_patches.py b/docker/scripts/apply_mineru_patches.py index 77ed4a7..f0e6e6f 100644 --- a/docker/scripts/apply_mineru_patches.py +++ b/docker/scripts/apply_mineru_patches.py @@ -15,9 +15,9 @@ 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) + 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') @@ -31,17 +31,17 @@ 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\))' + 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' + 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) + 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') @@ -52,8 +52,8 @@ def patch_b_predict_rec_batch(infer_dir): # 修改 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):', + r'( +)for rno in range\(len\(rec_result\)\):', + r'\1for rno in range(actual_batch_size):', c3 ) open(f, 'w').write(c4) @@ -63,9 +63,9 @@ 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) + 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')