paddle_table_cls.py 2.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768
  1. import os
  2. import cv2
  3. import numpy as np
  4. import onnxruntime
  5. from mineru.backend.pipeline.model_list import AtomicModel
  6. from mineru.utils.enum_class import ModelPath
  7. from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
  8. class PaddleTableClsModel:
  9. def __init__(self):
  10. self.sess = onnxruntime.InferenceSession(
  11. os.path.join(auto_download_and_get_model_root_path(ModelPath.paddle_table_cls), ModelPath.paddle_table_cls)
  12. )
  13. self.less_length = 256
  14. self.cw, self.ch = 224, 224
  15. self.std = [0.229, 0.224, 0.225]
  16. self.scale = 0.00392156862745098
  17. self.mean = [0.485, 0.456, 0.406]
  18. self.labels = [AtomicModel.WiredTable, AtomicModel.WirelessTable]
  19. def preprocess(self, img):
  20. # PIL图像转cv2
  21. img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
  22. img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  23. # 放大图片,使其最短边长为256
  24. h, w = img.shape[:2]
  25. scale = 256 / min(h, w)
  26. h_resize = round(h * scale)
  27. w_resize = round(w * scale)
  28. img = cv2.resize(img, (w_resize, h_resize), interpolation=1)
  29. # 调整为224*224的正方形
  30. h, w = img.shape[:2]
  31. cw, ch = 224, 224
  32. x1 = max(0, (w - cw) // 2)
  33. y1 = max(0, (h - ch) // 2)
  34. x2 = min(w, x1 + cw)
  35. y2 = min(h, y1 + ch)
  36. if w < cw or h < ch:
  37. raise ValueError(
  38. f"Input image ({w}, {h}) smaller than the target size ({cw}, {ch})."
  39. )
  40. img = img[y1:y2, x1:x2, ...]
  41. # 正则化
  42. split_im = list(cv2.split(img))
  43. std = [0.229, 0.224, 0.225]
  44. scale = 0.00392156862745098
  45. mean = [0.485, 0.456, 0.406]
  46. alpha = [scale / std[i] for i in range(len(std))]
  47. beta = [-mean[i] / std[i] for i in range(len(std))]
  48. for c in range(img.shape[2]):
  49. split_im[c] = split_im[c].astype(np.float32)
  50. split_im[c] *= alpha[c]
  51. split_im[c] += beta[c]
  52. img = cv2.merge(split_im)
  53. # 5. 转换为 CHW 格式
  54. img = img.transpose((2, 0, 1))
  55. imgs = [img]
  56. x = np.stack(imgs, axis=0).astype(dtype=np.float32, copy=False)
  57. return x
  58. def predict(self, img):
  59. x = self.preprocess(img)
  60. (result,) = self.sess.run(None, {"x": x})
  61. label = self.labels[np.argmax(result)]
  62. return label