Unimernet.py 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. import argparse
  2. import os
  3. import re
  4. import torch
  5. import unimernet.tasks as tasks
  6. from torch.utils.data import DataLoader, Dataset
  7. from torchvision import transforms
  8. from unimernet.common.config import Config
  9. from unimernet.processors import load_processor
  10. class MathDataset(Dataset):
  11. def __init__(self, image_paths, transform=None):
  12. self.image_paths = image_paths
  13. self.transform = transform
  14. def __len__(self):
  15. return len(self.image_paths)
  16. def latex_rm_whitespace(s: str):
  17. """Remove unnecessary whitespace from LaTeX code."""
  18. text_reg = r"(\\(operatorname|mathrm|text|mathbf)\s?\*? {.*?})"
  19. letter = "[a-zA-Z]"
  20. noletter = "[\W_^\d]"
  21. names = [x[0].replace(" ", "") for x in re.findall(text_reg, s)]
  22. s = re.sub(text_reg, lambda match: str(names.pop(0)), s)
  23. news = s
  24. while True:
  25. s = news
  26. news = re.sub(r"(?!\\ )(%s)\s+?(%s)" % (noletter, noletter), r"\1\2", s)
  27. news = re.sub(r"(?!\\ )(%s)\s+?(%s)" % (noletter, letter), r"\1\2", news)
  28. news = re.sub(r"(%s)\s+?(%s)" % (letter, noletter), r"\1\2", news)
  29. if news == s:
  30. break
  31. return s
  32. class UnimernetModel(object):
  33. def __init__(self, weight_dir, cfg_path, _device_="cpu"):
  34. args = argparse.Namespace(cfg_path=cfg_path, options=None)
  35. cfg = Config(args)
  36. cfg.config.model.pretrained = os.path.join(weight_dir, "pytorch_model.pth")
  37. cfg.config.model.model_config.model_name = weight_dir
  38. cfg.config.model.tokenizer_config.path = weight_dir
  39. task = tasks.setup_task(cfg)
  40. self.model = task.build_model(cfg)
  41. self.device = _device_
  42. self.model.to(_device_)
  43. self.model.eval()
  44. vis_processor = load_processor(
  45. "formula_image_eval",
  46. cfg.config.datasets.formula_rec_eval.vis_processor.eval,
  47. )
  48. self.mfr_transform = transforms.Compose(
  49. [
  50. vis_processor,
  51. ]
  52. )
  53. def predict(self, mfd_res, image):
  54. formula_list = []
  55. mf_image_list = []
  56. for xyxy, conf, cla in zip(
  57. mfd_res.boxes.xyxy.cpu(), mfd_res.boxes.conf.cpu(), mfd_res.boxes.cls.cpu()
  58. ):
  59. xmin, ymin, xmax, ymax = [int(p.item()) for p in xyxy]
  60. new_item = {
  61. "category_id": 13 + int(cla.item()),
  62. "poly": [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax],
  63. "score": round(float(conf.item()), 2),
  64. "latex": "",
  65. }
  66. formula_list.append(new_item)
  67. bbox_img = image[ymin:ymax, xmin:xmax]
  68. mf_image_list.append(bbox_img)
  69. dataset = MathDataset(mf_image_list, transform=self.mfr_transform)
  70. dataloader = DataLoader(dataset, batch_size=32, num_workers=0)
  71. mfr_res = []
  72. for mf_img in dataloader:
  73. mf_img = mf_img.to(self.device)
  74. with torch.no_grad():
  75. output = self.model.generate({"image": mf_img})
  76. mfr_res.extend(output["pred_str"])
  77. for res, latex in zip(formula_list, mfr_res):
  78. res["latex"] = latex_rm_whitespace(latex)
  79. return formula_list
  80. def batch_predict(self, images_mfd_res: list, images: list, batch_size: int = 64) -> list:
  81. images_formula_list = []
  82. mf_image_list = []
  83. backfill_list = []
  84. image_info = [] # Store (area, original_index, image) tuples
  85. # Collect images with their original indices
  86. for image_index in range(len(images_mfd_res)):
  87. mfd_res = images_mfd_res[image_index]
  88. np_array_image = images[image_index]
  89. formula_list = []
  90. for idx, (xyxy, conf, cla) in enumerate(zip(
  91. mfd_res.boxes.xyxy, mfd_res.boxes.conf, mfd_res.boxes.cls
  92. )):
  93. xmin, ymin, xmax, ymax = [int(p.item()) for p in xyxy]
  94. new_item = {
  95. "category_id": 13 + int(cla.item()),
  96. "poly": [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax],
  97. "score": round(float(conf.item()), 2),
  98. "latex": "",
  99. }
  100. formula_list.append(new_item)
  101. bbox_img = np_array_image[ymin:ymax, xmin:xmax]
  102. area = (xmax - xmin) * (ymax - ymin)
  103. curr_idx = len(mf_image_list)
  104. image_info.append((area, curr_idx, bbox_img))
  105. mf_image_list.append(bbox_img)
  106. images_formula_list.append(formula_list)
  107. backfill_list += formula_list
  108. # Stable sort by area
  109. image_info.sort(key=lambda x: x[0]) # sort by area
  110. sorted_indices = [x[1] for x in image_info]
  111. sorted_images = [x[2] for x in image_info]
  112. # Create mapping for results
  113. index_mapping = {new_idx: old_idx for new_idx, old_idx in enumerate(sorted_indices)}
  114. # Create dataset with sorted images
  115. dataset = MathDataset(sorted_images, transform=self.mfr_transform)
  116. dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=0)
  117. # Process batches and store results
  118. mfr_res = []
  119. for mf_img in dataloader:
  120. mf_img = mf_img.to(self.device)
  121. with torch.no_grad():
  122. output = self.model.generate({"image": mf_img})
  123. mfr_res.extend(output["pred_str"])
  124. # Restore original order
  125. unsorted_results = [""] * len(mfr_res)
  126. for new_idx, latex in enumerate(mfr_res):
  127. original_idx = index_mapping[new_idx]
  128. unsorted_results[original_idx] = latex_rm_whitespace(latex)
  129. # Fill results back
  130. for res, latex in zip(backfill_list, unsorted_results):
  131. res["latex"] = latex
  132. return images_formula_list