processors.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643
  1. import json
  2. import numpy as np
  3. import cv2
  4. import math
  5. import re
  6. from PIL import Image, ImageOps
  7. from typing import List, Optional, Tuple, Union, Dict, Any
  8. from loguru import logger
  9. from tokenizers import AddedToken
  10. from tokenizers import Tokenizer as TokenizerFast
  11. from mineru.model.mfr.unimernet.unimernet_hf.modeling_unimernet import fix_latex_left_right
  12. class UniMERNetImgDecode(object):
  13. """Class for decoding images for UniMERNet, including cropping margins, resizing, and padding."""
  14. def __init__(
  15. self, input_size: Tuple[int, int], random_padding: bool = False, **kwargs
  16. ) -> None:
  17. """Initializes the UniMERNetImgDecode class with input size and random padding options.
  18. Args:
  19. input_size (tuple): The desired input size for the images (height, width).
  20. random_padding (bool): Whether to use random padding for resizing.
  21. **kwargs: Additional keyword arguments."""
  22. self.input_size = input_size
  23. self.random_padding = random_padding
  24. def crop_margin(self, img: Image.Image) -> Image.Image:
  25. """Crops the margin of the image based on grayscale thresholding.
  26. Args:
  27. img (PIL.Image.Image): The input image.
  28. Returns:
  29. PIL.Image.Image: The cropped image."""
  30. data = np.array(img.convert("L"))
  31. data = data.astype(np.uint8)
  32. max_val = data.max()
  33. min_val = data.min()
  34. if max_val == min_val:
  35. return img
  36. data = (data - min_val) / (max_val - min_val) * 255
  37. gray = 255 * (data < 200).astype(np.uint8)
  38. coords = cv2.findNonZero(gray) # Find all non-zero points (text)
  39. a, b, w, h = cv2.boundingRect(coords) # Find minimum spanning bounding box
  40. return img.crop((a, b, w + a, h + b))
  41. def get_dimensions(self, img: Union[Image.Image, np.ndarray]) -> List[int]:
  42. """Gets the dimensions of the image.
  43. Args:
  44. img (PIL.Image.Image or numpy.ndarray): The input image.
  45. Returns:
  46. list: A list containing the number of channels, height, and width."""
  47. if hasattr(img, "getbands"):
  48. channels = len(img.getbands())
  49. else:
  50. channels = img.channels
  51. width, height = img.size
  52. return [channels, height, width]
  53. def _compute_resized_output_size(
  54. self,
  55. image_size: Tuple[int, int],
  56. size: Union[int, Tuple[int, int]],
  57. max_size: Optional[int] = None,
  58. ) -> List[int]:
  59. """Computes the resized output size of the image.
  60. Args:
  61. image_size (tuple): The original size of the image (height, width).
  62. size (int or tuple): The desired size for the smallest edge or both height and width.
  63. max_size (int, optional): The maximum allowed size for the longer edge.
  64. Returns:
  65. list: A list containing the new height and width."""
  66. if len(size) == 1: # specified size only for the smallest edge
  67. h, w = image_size
  68. short, long = (w, h) if w <= h else (h, w)
  69. requested_new_short = size if isinstance(size, int) else size[0]
  70. new_short, new_long = requested_new_short, int(
  71. requested_new_short * long / short
  72. )
  73. if max_size is not None:
  74. if max_size <= requested_new_short:
  75. raise ValueError(
  76. f"max_size = {max_size} must be strictly greater than the requested "
  77. f"size for the smaller edge size = {size}"
  78. )
  79. if new_long > max_size:
  80. new_short, new_long = int(max_size * new_short / new_long), max_size
  81. new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short)
  82. else: # specified both h and w
  83. new_w, new_h = size[1], size[0]
  84. return [new_h, new_w]
  85. def resize(
  86. self, img: Image.Image, size: Union[int, Tuple[int, int]]
  87. ) -> Image.Image:
  88. """Resizes the image to the specified size.
  89. Args:
  90. img (PIL.Image.Image): The input image.
  91. size (int or tuple): The desired size for the smallest edge or both height and width.
  92. Returns:
  93. PIL.Image.Image: The resized image."""
  94. _, image_height, image_width = self.get_dimensions(img)
  95. if isinstance(size, int):
  96. size = [size]
  97. max_size = None
  98. output_size = self._compute_resized_output_size(
  99. (image_height, image_width), size, max_size
  100. )
  101. img = img.resize(tuple(output_size[::-1]), resample=2)
  102. return img
  103. def img_decode(self, img: np.ndarray) -> Optional[np.ndarray]:
  104. """Decodes the image by cropping margins, resizing, and adding padding.
  105. Args:
  106. img (numpy.ndarray): The input image array.
  107. Returns:
  108. numpy.ndarray: The decoded image array."""
  109. try:
  110. img = self.crop_margin(Image.fromarray(img).convert("RGB"))
  111. except OSError:
  112. return
  113. if img.height == 0 or img.width == 0:
  114. return
  115. img = self.resize(img, min(self.input_size))
  116. img.thumbnail((self.input_size[1], self.input_size[0]))
  117. delta_width = self.input_size[1] - img.width
  118. delta_height = self.input_size[0] - img.height
  119. if self.random_padding:
  120. pad_width = np.random.randint(low=0, high=delta_width + 1)
  121. pad_height = np.random.randint(low=0, high=delta_height + 1)
  122. else:
  123. pad_width = delta_width // 2
  124. pad_height = delta_height // 2
  125. padding = (
  126. pad_width,
  127. pad_height,
  128. delta_width - pad_width,
  129. delta_height - pad_height,
  130. )
  131. return np.array(ImageOps.expand(img, padding))
  132. def __call__(self, imgs: List[np.ndarray]) -> List[Optional[np.ndarray]]:
  133. """Calls the img_decode method on a list of images.
  134. Args:
  135. imgs (list of numpy.ndarray): The list of input image arrays.
  136. Returns:
  137. list of numpy.ndarray: The list of decoded image arrays."""
  138. return [self.img_decode(img) for img in imgs]
  139. class UniMERNetTestTransform:
  140. """
  141. A class for transforming images according to UniMERNet test specifications.
  142. """
  143. def __init__(self, **kwargs) -> None:
  144. """
  145. Initializes the UniMERNetTestTransform class.
  146. """
  147. super().__init__()
  148. self.num_output_channels = 3
  149. def transform(self, img: np.ndarray) -> np.ndarray:
  150. """
  151. Transforms a single image for UniMERNet testing.
  152. Args:
  153. img (numpy.ndarray): The input image.
  154. Returns:
  155. numpy.ndarray: The transformed image.
  156. """
  157. mean = [0.7931, 0.7931, 0.7931]
  158. std = [0.1738, 0.1738, 0.1738]
  159. scale = float(1 / 255.0)
  160. shape = (1, 1, 3)
  161. mean = np.array(mean).reshape(shape).astype("float32")
  162. std = np.array(std).reshape(shape).astype("float32")
  163. img = (img.astype("float32") * scale - mean) / std
  164. grayscale_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  165. squeezed = np.squeeze(grayscale_image)
  166. img = cv2.merge([squeezed] * self.num_output_channels)
  167. return img
  168. def __call__(self, imgs: List[np.ndarray]) -> List[np.ndarray]:
  169. """
  170. Applies the transform to a list of images.
  171. Args:
  172. imgs (list of numpy.ndarray): The list of input images.
  173. Returns:
  174. list of numpy.ndarray: The list of transformed images.
  175. """
  176. return [self.transform(img) for img in imgs]
  177. class LatexImageFormat:
  178. """Class for formatting images to a specific format suitable for LaTeX."""
  179. def __init__(self, **kwargs) -> None:
  180. """Initializes the LatexImageFormat class with optional keyword arguments."""
  181. super().__init__()
  182. def format(self, img: np.ndarray) -> np.ndarray:
  183. """Formats a single image to the LaTeX-compatible format.
  184. Args:
  185. img (numpy.ndarray): The input image as a numpy array.
  186. Returns:
  187. numpy.ndarray: The formatted image as a numpy array with an added dimension for color.
  188. """
  189. im_h, im_w = img.shape[:2]
  190. divide_h = math.ceil(im_h / 16) * 16
  191. divide_w = math.ceil(im_w / 16) * 16
  192. img = img[:, :, 0]
  193. img = np.pad(
  194. img, ((0, divide_h - im_h), (0, divide_w - im_w)), constant_values=(1, 1)
  195. )
  196. img_expanded = img[:, :, np.newaxis].transpose(2, 0, 1)
  197. return img_expanded[np.newaxis, :]
  198. def __call__(self, imgs: List[np.ndarray]) -> List[np.ndarray]:
  199. """Applies the format method to a list of images.
  200. Args:
  201. imgs (list of numpy.ndarray): A list of input images as numpy arrays.
  202. Returns:
  203. list of numpy.ndarray: A list of formatted images as numpy arrays.
  204. """
  205. return [self.format(img) for img in imgs]
  206. class ToBatch(object):
  207. """A class for batching images."""
  208. def __init__(self, **kwargs) -> None:
  209. """Initializes the ToBatch object."""
  210. super(ToBatch, self).__init__()
  211. def __call__(self, imgs: List[np.ndarray]) -> List[np.ndarray]:
  212. """Concatenates a list of images into a single batch.
  213. Args:
  214. imgs (list): A list of image arrays to be concatenated.
  215. Returns:
  216. list: A list containing the concatenated batch of images wrapped in another list (to comply with common batch processing formats).
  217. """
  218. batch_imgs = np.concatenate(imgs)
  219. batch_imgs = batch_imgs.copy()
  220. x = [batch_imgs]
  221. return x
  222. class UniMERNetDecode(object):
  223. """Class for decoding tokenized inputs using UniMERNet tokenizer.
  224. Attributes:
  225. SPECIAL_TOKENS_ATTRIBUTES (List[str]): List of special token attributes.
  226. model_input_names (List[str]): List of model input names.
  227. max_seq_len (int): Maximum sequence length.
  228. pad_token_id (int): ID for the padding token.
  229. bos_token_id (int): ID for the beginning-of-sequence token.
  230. eos_token_id (int): ID for the end-of-sequence token.
  231. padding_side (str): Padding side, either 'left' or 'right'.
  232. pad_token (str): Padding token.
  233. pad_token_type_id (int): Type ID for the padding token.
  234. pad_to_multiple_of (Optional[int]): If set, pad to a multiple of this value.
  235. tokenizer (TokenizerFast): Fast tokenizer instance.
  236. Args:
  237. character_list (Dict[str, Any]): Dictionary containing tokenizer configuration.
  238. **kwargs: Additional keyword arguments.
  239. """
  240. SPECIAL_TOKENS_ATTRIBUTES = [
  241. "bos_token",
  242. "eos_token",
  243. "unk_token",
  244. "sep_token",
  245. "pad_token",
  246. "cls_token",
  247. "mask_token",
  248. "additional_special_tokens",
  249. ]
  250. def __init__(
  251. self,
  252. character_list: Dict[str, Any],
  253. **kwargs,
  254. ) -> None:
  255. """Initializes the UniMERNetDecode class.
  256. Args:
  257. character_list (Dict[str, Any]): Dictionary containing tokenizer configuration.
  258. **kwargs: Additional keyword arguments.
  259. """
  260. self._unk_token = "<unk>"
  261. self._bos_token = "<s>"
  262. self._eos_token = "</s>"
  263. self._pad_token = "<pad>"
  264. self._sep_token = None
  265. self._cls_token = None
  266. self._mask_token = None
  267. self._additional_special_tokens = []
  268. self.model_input_names = ["input_ids", "token_type_ids", "attention_mask"]
  269. self.max_seq_len = 2048
  270. self.pad_token_id = 1
  271. self.bos_token_id = 0
  272. self.eos_token_id = 2
  273. self.padding_side = "right"
  274. self.pad_token_id = 1
  275. self.pad_token = "<pad>"
  276. self.pad_token_type_id = 0
  277. self.pad_to_multiple_of = None
  278. fast_tokenizer_str = json.dumps(character_list["fast_tokenizer_file"])
  279. fast_tokenizer_buffer = fast_tokenizer_str.encode("utf-8")
  280. self.tokenizer = TokenizerFast.from_buffer(fast_tokenizer_buffer)
  281. tokenizer_config = (
  282. character_list["tokenizer_config_file"]
  283. if "tokenizer_config_file" in character_list
  284. else None
  285. )
  286. added_tokens_decoder = {}
  287. added_tokens_map = {}
  288. if tokenizer_config is not None:
  289. init_kwargs = tokenizer_config
  290. if "added_tokens_decoder" in init_kwargs:
  291. for idx, token in init_kwargs["added_tokens_decoder"].items():
  292. if isinstance(token, dict):
  293. token = AddedToken(**token)
  294. if isinstance(token, AddedToken):
  295. added_tokens_decoder[int(idx)] = token
  296. added_tokens_map[str(token)] = token
  297. else:
  298. raise ValueError(
  299. f"Found a {token.__class__} in the saved `added_tokens_decoder`, should be a dictionary or an AddedToken instance"
  300. )
  301. init_kwargs["added_tokens_decoder"] = added_tokens_decoder
  302. added_tokens_decoder = init_kwargs.pop("added_tokens_decoder", {})
  303. tokens_to_add = [
  304. token
  305. for index, token in sorted(
  306. added_tokens_decoder.items(), key=lambda x: x[0]
  307. )
  308. if token not in added_tokens_decoder
  309. ]
  310. added_tokens_encoder = self.added_tokens_encoder(added_tokens_decoder)
  311. encoder = list(added_tokens_encoder.keys()) + [
  312. str(token) for token in tokens_to_add
  313. ]
  314. tokens_to_add += [
  315. token
  316. for token in self.all_special_tokens_extended
  317. if token not in encoder and token not in tokens_to_add
  318. ]
  319. if len(tokens_to_add) > 0:
  320. is_last_special = None
  321. tokens = []
  322. special_tokens = self.all_special_tokens
  323. for token in tokens_to_add:
  324. is_special = (
  325. (token.special or str(token) in special_tokens)
  326. if isinstance(token, AddedToken)
  327. else str(token) in special_tokens
  328. )
  329. if is_last_special is None or is_last_special == is_special:
  330. tokens.append(token)
  331. else:
  332. self._add_tokens(tokens, special_tokens=is_last_special)
  333. tokens = [token]
  334. is_last_special = is_special
  335. if tokens:
  336. self._add_tokens(tokens, special_tokens=is_last_special)
  337. def _add_tokens(
  338. self, new_tokens: "List[Union[AddedToken, str]]", special_tokens: bool = False
  339. ) -> "List[Union[AddedToken, str]]":
  340. """Adds new tokens to the tokenizer.
  341. Args:
  342. new_tokens (List[Union[AddedToken, str]]): Tokens to be added.
  343. special_tokens (bool): Indicates whether the tokens are special tokens.
  344. Returns:
  345. List[Union[AddedToken, str]]: added tokens.
  346. """
  347. if special_tokens:
  348. return self.tokenizer.add_special_tokens(new_tokens)
  349. return self.tokenizer.add_tokens(new_tokens)
  350. def added_tokens_encoder(
  351. self, added_tokens_decoder: "Dict[int, AddedToken]"
  352. ) -> Dict[str, int]:
  353. """Creates an encoder dictionary from added tokens.
  354. Args:
  355. added_tokens_decoder (Dict[int, AddedToken]): Dictionary mapping token IDs to tokens.
  356. Returns:
  357. Dict[str, int]: Dictionary mapping token strings to IDs.
  358. """
  359. return {
  360. k.content: v
  361. for v, k in sorted(added_tokens_decoder.items(), key=lambda item: item[0])
  362. }
  363. @property
  364. def all_special_tokens(self) -> List[str]:
  365. """Retrieves all special tokens.
  366. Returns:
  367. List[str]: List of all special tokens as strings.
  368. """
  369. all_toks = [str(s) for s in self.all_special_tokens_extended]
  370. return all_toks
  371. @property
  372. def all_special_tokens_extended(self) -> "List[Union[str, AddedToken]]":
  373. """Retrieves all special tokens, including extended ones.
  374. Returns:
  375. List[Union[str, AddedToken]]: List of all special tokens.
  376. """
  377. all_tokens = []
  378. seen = set()
  379. for value in self.special_tokens_map_extended.values():
  380. if isinstance(value, (list, tuple)):
  381. tokens_to_add = [token for token in value if str(token) not in seen]
  382. else:
  383. tokens_to_add = [value] if str(value) not in seen else []
  384. seen.update(map(str, tokens_to_add))
  385. all_tokens.extend(tokens_to_add)
  386. return all_tokens
  387. @property
  388. def special_tokens_map_extended(self) -> Dict[str, Union[str, List[str]]]:
  389. """Retrieves the extended map of special tokens.
  390. Returns:
  391. Dict[str, Union[str, List[str]]]: Dictionary mapping special token attributes to their values.
  392. """
  393. set_attr = {}
  394. for attr in self.SPECIAL_TOKENS_ATTRIBUTES:
  395. attr_value = getattr(self, "_" + attr)
  396. if attr_value:
  397. set_attr[attr] = attr_value
  398. return set_attr
  399. def convert_ids_to_tokens(
  400. self, ids: Union[int, List[int]], skip_special_tokens: bool = False
  401. ) -> Union[str, List[str]]:
  402. """Converts token IDs to token strings.
  403. Args:
  404. ids (Union[int, List[int]]): Token ID(s) to convert.
  405. skip_special_tokens (bool): Whether to skip special tokens during conversion.
  406. Returns:
  407. Union[str, List[str]]: Converted token string(s).
  408. """
  409. if isinstance(ids, int):
  410. return self.tokenizer.id_to_token(ids)
  411. tokens = []
  412. for index in ids:
  413. index = int(index)
  414. if skip_special_tokens and index in self.all_special_ids:
  415. continue
  416. tokens.append(self.tokenizer.id_to_token(index))
  417. return tokens
  418. def detokenize(self, tokens: List[List[int]]) -> List[List[str]]:
  419. """Detokenizes a list of token IDs back into strings.
  420. Args:
  421. tokens (List[List[int]]): List of token ID lists.
  422. Returns:
  423. List[List[str]]: List of detokenized strings.
  424. """
  425. self.tokenizer.bos_token = "<s>"
  426. self.tokenizer.eos_token = "</s>"
  427. self.tokenizer.pad_token = "<pad>"
  428. toks = [self.convert_ids_to_tokens(tok) for tok in tokens]
  429. for b in range(len(toks)):
  430. for i in reversed(range(len(toks[b]))):
  431. if toks[b][i] is None:
  432. toks[b][i] = ""
  433. toks[b][i] = toks[b][i].replace("Ġ", " ").strip()
  434. if toks[b][i] in (
  435. [
  436. self.tokenizer.bos_token,
  437. self.tokenizer.eos_token,
  438. self.tokenizer.pad_token,
  439. ]
  440. ):
  441. del toks[b][i]
  442. return toks
  443. def token2str(self, token_ids: List[List[int]]) -> List[str]:
  444. """Converts a list of token IDs to strings.
  445. Args:
  446. token_ids (List[List[int]]): List of token ID lists.
  447. Returns:
  448. List[str]: List of converted strings.
  449. """
  450. generated_text = []
  451. for tok_id in token_ids:
  452. end_idx = np.argwhere(tok_id == 2)
  453. if len(end_idx) > 0:
  454. end_idx = int(end_idx[0][0])
  455. tok_id = tok_id[: end_idx + 1]
  456. generated_text.append(
  457. self.tokenizer.decode(tok_id, skip_special_tokens=True)
  458. )
  459. generated_text = [self.post_process(text) for text in generated_text]
  460. return generated_text
  461. def normalize(self, s: str) -> str:
  462. """Normalizes a string by removing unnecessary spaces.
  463. Args:
  464. s (str): String to normalize.
  465. Returns:
  466. str: Normalized string.
  467. """
  468. text_reg = r"(\\(operatorname|mathrm|text|mathbf)\s?\*? {.*?})"
  469. letter = "[a-zA-Z]"
  470. noletter = r"[\W_^\d]"
  471. names = []
  472. for x in re.findall(text_reg, s):
  473. pattern = r"(\\[a-zA-Z]+)\s(?=\w)|\\[a-zA-Z]+\s(?=})"
  474. matches = re.findall(pattern, x[0])
  475. for m in matches:
  476. if (
  477. m
  478. not in [
  479. "\\operatorname",
  480. "\\mathrm",
  481. "\\text",
  482. "\\mathbf",
  483. ]
  484. and m.strip() != ""
  485. ):
  486. s = s.replace(m, m + "XXXXXXX")
  487. s = s.replace(" ", "")
  488. names.append(s)
  489. if len(names) > 0:
  490. s = re.sub(text_reg, lambda match: str(names.pop(0)), s)
  491. news = s
  492. while True:
  493. s = news
  494. news = re.sub(r"(?!\\ )(%s)\s+?(%s)" % (noletter, noletter), r"\1\2", s)
  495. news = re.sub(r"(?!\\ )(%s)\s+?(%s)" % (noletter, letter), r"\1\2", news)
  496. news = re.sub(r"(%s)\s+?(%s)" % (letter, noletter), r"\1\2", news)
  497. if news == s:
  498. break
  499. return s.replace("XXXXXXX", " ")
  500. def remove_chinese_text_wrapping(self, formula):
  501. pattern = re.compile(r"\\text\s*{\s*([^}]*?[\u4e00-\u9fff]+[^}]*?)\s*}")
  502. def replacer(match):
  503. return match.group(1)
  504. replaced_formula = pattern.sub(replacer, formula)
  505. return replaced_formula.replace('"', "")
  506. UP_PATTERN = re.compile(r'\\up([a-zA-Z]+)')
  507. def post_process(self, text: str) -> str:
  508. """Post-processes a string by fixing text and normalizing it.
  509. Args:
  510. text (str): String to post-process.
  511. Returns:
  512. str: Post-processed string.
  513. """
  514. from ftfy import fix_text
  515. text = self.remove_chinese_text_wrapping(text)
  516. text = fix_text(text)
  517. text = fix_latex_left_right(text)
  518. text = self.UP_PATTERN.sub(
  519. lambda m: m.group(0) if m.group(1) in ["arrow", "downarrow", "lus", "silon"] else f"\\{m.group(1)}", text
  520. )
  521. text = self.normalize(text)
  522. return text
  523. def __call__(
  524. self,
  525. preds: np.ndarray,
  526. label: Optional[np.ndarray] = None,
  527. mode: str = "eval",
  528. *args,
  529. **kwargs,
  530. ) -> Union[List[str], tuple]:
  531. """Processes predictions and optionally labels, returning the decoded text.
  532. Args:
  533. preds (np.ndarray): Model predictions.
  534. label (Optional[np.ndarray]): True labels, if available.
  535. mode (str): Mode of operation, either 'train' or 'eval'.
  536. Returns:
  537. Union[List[str], tuple]: Decoded text, optionally with labels.
  538. """
  539. if mode == "train":
  540. preds_idx = np.array(preds.argmax(axis=2))
  541. text = self.token2str(preds_idx)
  542. else:
  543. text = self.token2str(np.array(preds))
  544. if label is None:
  545. return text
  546. label = self.token2str(np.array(label))
  547. return text, label