Unimernet.py 7.8 KB

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