from typing import Tuple, List, Sequence, Optional, Union from pathlib import Path import re import torch import tokenizers as tk from PIL import Image from matplotlib import pyplot as plt from matplotlib import patches from torchvision import transforms from torch import nn, Tensor from functools import partial import numpy.typing as npt from numpy import uint8 ImageType = npt.NDArray[uint8] import warnings import time import argparse from bs4 import BeautifulSoup as bs from .src.model import EncoderDecoder, ImgLinearBackbone, Encoder, Decoder from .src.utils import subsequent_mask, pred_token_within_range, greedy_sampling, bbox_str_to_token_list, html_str_to_token_list,cell_str_to_token_list, build_table_from_html_and_cell, html_table_template from .src.trainer.utils import VALID_HTML_TOKEN, VALID_BBOX_TOKEN, INVALID_CELL_TOKEN warnings.filterwarnings('ignore') class UnitableFullSinglePredictor(): def __init__(self): MODEL_FILE_NAME = ["unitable_large_structure.pt", "unitable_large_bbox.pt", "unitable_large_content.pt"] MODEL_DIR = Path("unitable/experiments/unitable_weights") # UniTable large model self.d_model = 768 self.patch_size = 16 self.nhead = 12 self.dropout = 0.2 self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.backbone= ImgLinearBackbone(d_model=self.d_model, patch_size=self.patch_size) self.encoder= Encoder( d_model=self.d_model, nhead=self.nhead, dropout=self.dropout, activation="gelu", norm_first=True, nlayer=12, ff_ratio=4, ) self.decoder= Decoder( d_model=self.d_model, nhead=self.nhead, dropout=self.dropout, activation="gelu", norm_first=True, nlayer=4, ff_ratio=4, ) """ start1 = time.time() # Table structure extraction self.vocabS, self.modelS = self.load_vocab_and_model( backbone= ImgLinearBackbone(d_model=self.d_model, patch_size=self.patch_size), encoder= Encoder( d_model=self.d_model, nhead=self.nhead, dropout=self.dropout, activation="gelu", norm_first=True, nlayer=12, ff_ratio=4, ), decoder= Decoder( d_model=self.d_model, nhead=self.nhead, dropout=self.dropout, activation="gelu", norm_first=True, nlayer=4, ff_ratio=4, ), d_model= self.d_model, dropout= self.dropout, vocab_path="unitable/vocab/vocab_html.json", max_seq_len=784, model_weights=MODEL_DIR / MODEL_FILE_NAME[0] ) end1 = time.time() print("time to load table structure model ",end1-start1,"seconds") start3 = time.time() # Table cell bbox detection self.vocabB, self.modelB = self.load_vocab_and_model( backbone = ImgLinearBackbone(d_model=self.d_model, patch_size=self.patch_size), encoder = Encoder( d_model= self.d_model, nhead= self.nhead, dropout = self.dropout, activation="gelu", norm_first=True, nlayer=12, ff_ratio=4, ), decoder = Decoder( d_model= self.d_model, nhead= self.nhead, dropout = self.dropout, activation="gelu", norm_first=True, nlayer=4, ff_ratio=4, ), d_model= self.d_model, dropout= self.dropout, vocab_path="unitable/vocab/vocab_bbox.json", max_seq_len=1024, model_weights=MODEL_DIR / MODEL_FILE_NAME[1], ) end3 = time.time() print("time to load cell bbox detection model ",end3-start3,"seconds") start4 = time.time() # Table cell bbox detection self.vocabC, self.modelC = self.load_vocab_and_model( backbone = ImgLinearBackbone(d_model=self.d_model, patch_size=self.patch_size), encoder = Encoder( d_model= self.d_model, nhead= self.nhead, dropout = self.dropout, activation="gelu", norm_first=True, nlayer=12, ff_ratio=4, ), decoder = Decoder( d_model= self.d_model, nhead= self.nhead, dropout = self.dropout, activation="gelu", norm_first=True, nlayer=4, ff_ratio=4, ), d_model= self.d_model, dropout= self.dropout, vocab_path="unitable/vocab/vocab_cell_6k.json", max_seq_len=200, #Using the content recognition model i guess model_weights=MODEL_DIR / MODEL_FILE_NAME[2], ) end4 = time.time() print("time to load cell recognition model ",end4-start4,"seconds") """ def load_vocab_and_model( self, vocab_path: Union[str, Path], max_seq_len: int, model_weights: Union[str, Path], ) -> Tuple[tk.Tokenizer, EncoderDecoder]: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") vocab = tk.Tokenizer.from_file(vocab_path) model = EncoderDecoder( backbone= self.backbone, encoder= self.encoder, decoder= self.decoder, vocab_size= vocab.get_vocab_size(), d_model= self.d_model, padding_idx= vocab.token_to_id(""), max_seq_len=max_seq_len, dropout=self.dropout, norm_layer=partial(nn.LayerNorm, eps=1e-6) ) # it loads weights onto the CPU first and then moves the model to the desired device model.load_state_dict(torch.load(model_weights, map_location="cpu")) model = model.to(device) return vocab, model def autoregressive_decode( self, model: EncoderDecoder, image: Tensor, prefix: Sequence[int], max_decode_len: int, eos_id: int, token_whitelist: Optional[Sequence[int]] = None, token_blacklist: Optional[Sequence[int]] = None, ) -> Tensor: model.eval() with torch.no_grad(): memory = model.encode(image) context = torch.tensor(prefix, dtype=torch.int32).repeat(image.shape[0], 1).to(self.device) for _ in range(max_decode_len): eos_flag = [eos_id in k for k in context] if all(eos_flag): break with torch.no_grad(): causal_mask = subsequent_mask(context.shape[1]).to(self.device) logits = model.decode( memory, context, tgt_mask=causal_mask, tgt_padding_mask=None ) logits = model.generator(logits)[:, -1, :] logits = pred_token_within_range( logits.detach(), white_list=token_whitelist, black_list=token_blacklist, ) next_probs, next_tokens = greedy_sampling(logits) context = torch.cat([context, next_tokens], dim=1) return context @staticmethod def image_to_tensor(image: Image, size: Tuple[int, int]) -> Tensor: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Resize the image with padding #resized_image = UnitableFullPredictor.resize_with_padding(image, size) T = transforms.Compose([ transforms.Resize(size), transforms.ToTensor(), transforms.Normalize(mean=[0.86597056,0.88463002,0.87491087], std = [0.20686628,0.18201602,0.18485524]) ]) image_tensor = T(image) image_tensor = image_tensor.to(device).unsqueeze(0) return image_tensor """ @staticmethod def resize_with_padding(image: Image, target_size: Tuple[int, int]) -> Image: #Resize the image to fit within the target size while preserving aspect ratio, #then add padding to match the target size. original_width, original_height = image.size target_width, target_height = target_size # Calculate the new size preserving aspect ratio aspect_ratio = original_width / original_height if target_width / target_height > aspect_ratio: new_height = target_height new_width = int(new_height * aspect_ratio) else: new_width = target_width new_height = int(new_width / aspect_ratio) # Resize the image to the new size resized_image = image.resize((new_width, new_height),Image.LANCZOS) # Create a new image with white background new_image = Image.new("RGB", (target_width, target_height), (255, 255, 255)) # Paste the resized image onto the white background paste_position = ((target_width - new_width) // 2, (target_height - new_height) // 2) new_image.paste(resized_image, paste_position) new_image.save("../res/table_resize_with_padding.png") return new_image """ def rescale_bbox( self, bbox: Sequence[Sequence[float]], src: Tuple[int, int], tgt: Tuple[int, int] ) -> Sequence[Sequence[float]]: assert len(src) == len(tgt) == 2 ratio = [tgt[0] / src[0], tgt[1] / src[1]] * 2 print(ratio) bbox = [[int(round(i * j)) for i, j in zip(entry, ratio)] for entry in bbox] return bbox """ @staticmethod def rescale_bbox( bbox: Sequence[Sequence[float]], src: Tuple[int, int], tgt: Tuple[int, int] ) -> Sequence[Sequence[float]]: #Rescale bounding boxes according to the transformation applied in resize_with_padding. src_width, src_height = src tgt_width, tgt_height = tgt # Calculate the new size preserving aspect ratio aspect_ratio = src_width / src_height if tgt_width / tgt_height > aspect_ratio: new_height = tgt_height new_width = int(new_height * aspect_ratio) else: new_width = tgt_width new_height = int(new_width / aspect_ratio) # Calculate the scale factors #THIS *2 factor was done in their code - why ? i have no clue scale_x = (new_width / src_width ) * 2 scale_y = (new_height / src_height) *2 # Calculate the padding pad_x = (tgt_width - new_width) // 2 pad_y = (tgt_height - new_height) // 2 # Rescale and adjust the bounding boxes rescaled_bbox = [] for entry in bbox: x_min = int(round(entry[0] * scale_x -pad_x)) y_min = int(round(entry[1] * scale_y - pad_y)) x_max = int(round(entry[2] * scale_x - pad_x)) y_max = int(round(entry[3] * scale_y - pad_y)) rescaled_bbox.append([x_min, y_min, x_max, y_max]) return rescaled_bbox """ def predict(self, image:ImageType): MODEL_FILE_NAME = ["unitable_large_structure.pt", "unitable_large_bbox.pt", "unitable_large_content.pt"] MODEL_DIR = Path("unitable/experiments/unitable_weights") image_size = image.size print("RUNING SINGLE IMAGE UNITABLE FOR DEBUGGGING ") # Image transformation image_tensor = self.image_to_tensor(image, (448, 448)) #print(image_tensor) """ Step 1 Table Structure recognition """ start1 = time.time() # Table structure extraction vocabS, modelS = self.load_vocab_and_model( vocab_path="unitable/vocab/vocab_html.json", max_seq_len=784, model_weights=MODEL_DIR / MODEL_FILE_NAME[0] ) end1 = time.time() print("time to load table structure model ",end1-start1,"seconds") start2 = time.time() # Inference pred_html = self.autoregressive_decode( model= modelS, image= image_tensor, prefix=[vocabS.token_to_id("[html]")], max_decode_len=512, eos_id=vocabS.token_to_id(""), token_whitelist=[vocabS.token_to_id(i) for i in VALID_HTML_TOKEN], token_blacklist = None ) end2 = time.time() print("time for inference table structure ",end2-start2,"seconds") # Convert token id to token text pred_html = pred_html.detach().cpu().numpy()[0] pred_html = vocabS.decode(pred_html, skip_special_tokens=False) #print(pred_html) pred_html = html_str_to_token_list(pred_html) print(pred_html) """ Step 2 Table Cell detection """ start3 = time.time() # Table cell bbox detection vocabB, modelB = self.load_vocab_and_model( vocab_path="unitable/vocab/vocab_bbox.json", max_seq_len=1024, model_weights=MODEL_DIR / MODEL_FILE_NAME[1], ) end3 = time.time() print("time to load cell bbox detection model ",end3-start3,"seconds") start4 = time.time() # Inference pred_bbox = self.autoregressive_decode( model=modelB, image=image_tensor, prefix=[vocabB.token_to_id("[bbox]")], max_decode_len=1024, eos_id=vocabB.token_to_id(""), token_whitelist=[vocabB.token_to_id(i) for i in VALID_BBOX_TOKEN[: 449]], token_blacklist = None ) end4 = time.time() print("time to do inference for table cell bbox detection model ",end4-start4,"seconds") # Convert token id to token text pred_bbox = pred_bbox.detach().cpu().numpy()[0] pred_bbox = vocabB.decode(pred_bbox, skip_special_tokens=False) pred_bbox = bbox_str_to_token_list(pred_bbox) pred_bbox = self.rescale_bbox(pred_bbox, src=(448, 448), tgt=image.size) print(pred_bbox) print("Size of the image ") #(1498, 971) print(image.size) print("Number of bounding boxes ") print(len(pred_bbox)) countcells = 0 #startBody = False #startFirstRow = True #numElemInRow = 0 for elem in pred_html : #if elem == '': # startBody = True #elif startBody ==True and elem == '': # startFirstRow = True #elif startFirstRow == True and elem == '[]': # numElemInRow +=1 #elif startBody ==True and elem == '': # startFirstRow = False # startBody = False if elem == '[]': countcells+=1 #275 print(countcells) if countcells > len(pred_bbox): #TODO Extra processing for big tables #Find the last incomplete row and its ymax coordinate # Last bbox's ymax gives us coordinate of where the cutted off row starts #IMPORTANT : pred_bbox is xmin, ymax, xmax, ymin cut_off = pred_bbox[-1][1] width = image.size[0] height = image.size[1] #bbox = (0, cut_off, width, height) #IMPORTANT : crop takes in (xmin, ymax, xmax, ymin) coordintes !!! bbox = (0, cut_off, width, height) # Crop the image to the specified bounding box cropped_image = image.crop(bbox) cropped_image.save("./res/cropped_image_for_extra_bbox_det.png") image_tensor = self.image_to_tensor(cropped_image, (448, 448)) pred_bbox_extra = self.autoregressive_decode( model=modelB, image=image_tensor, prefix=[vocabB.token_to_id("[bbox]")], max_decode_len=1024, eos_id=vocabB.token_to_id(""), token_whitelist=[vocabB.token_to_id(i) for i in VALID_BBOX_TOKEN[: 449]], token_blacklist = None ) # Convert token id to token text pred_bbox_extra = pred_bbox_extra.detach().cpu().numpy()[0] pred_bbox_extra = vocabB.decode(pred_bbox_extra, skip_special_tokens=False) pred_bbox_extra = bbox_str_to_token_list(pred_bbox_extra) numberOrCellsToAdd = countcells-len(pred_bbox) pred_bbox_extra = pred_bbox_extra[-numberOrCellsToAdd:] pred_bbox_extra = self.rescale_bbox(pred_bbox_extra, src=(448, 448), tgt=cropped_image.size) #This resulted in table_bbox_test_extra_3.png #pred_bbox_extra = [[i[0], i[1]+cut_off, i[2], i[3]+cut_off] for i in pred_bbox_extra] pred_bbox_extra = [[i[0], i[1]+cut_off, i[2], i[3]+cut_off] for i in pred_bbox_extra] pred_bbox = pred_bbox + pred_bbox_extra #[[25, 63, 152, 86], [227, 63, 292, 86], [326, 63, 373, 86], [413, 63, 460, 86], [562, 63, 609, 86], [708, 63, 758, 86], [848, 63, 895, 86], [935, 63, 982, 86], [1025, 63, 1075, 86], [1119, 63, 1165, 86], [1280, 63, 1327, 86]] print(pred_bbox_extra) #11 print(len(pred_bbox_extra)) fig, ax = plt.subplots(figsize=(12, 10)) for i in pred_bbox: #i is xmin, ymin, xmax, ymax based on the function usage rect = patches.Rectangle(i[:2], i[2] - i[0], i[3] - i[1], linewidth=1, edgecolor='r', facecolor='none') ax.add_patch(rect) ax.set_axis_off() ax.imshow(image) fig.savefig('./res/table_debug3/singleimageres.png', bbox_inches='tight', dpi=300) """ Step 3 : Table cell content recognition """ start4 = time.time() # Table cell bbox detection vocabC, modelC = self.load_vocab_and_model( vocab_path="unitable/vocab/vocab_cell_6k.json", max_seq_len=200, model_weights=MODEL_DIR / MODEL_FILE_NAME[2], ) end4 = time.time() print("time to load cell recognition model ",end4-start4,"seconds") # Cell image cropping and transformation """ images = [image.crop(bbox) for bbox in pred_bbox] for idx, img in enumerate(images): img.save("res/debug/cell_{}.png".format(idx)) """ #Cropping boundaries are fine image_tensor = [self.image_to_tensor(image.crop(bbox), size=(112, 448)) for bbox in pred_bbox] image_tensor = torch.cat(image_tensor, dim=0) #print("size of tensor") #print(image_tensor.size()) start4 = time.time() # Inference pred_cell = self.autoregressive_decode( model=modelC, image=image_tensor, prefix=[vocabC.token_to_id("[cell]")], max_decode_len=200, eos_id=vocabC.token_to_id(""), token_whitelist=None, token_blacklist = [vocabC.token_to_id(i) for i in INVALID_CELL_TOKEN] ) # Convert token id to token text pred_cell = pred_cell.detach().cpu().numpy() pred_cell = vocabC.decode_batch(pred_cell, skip_special_tokens=False) end4 = time.time() print("time to do cell recognition ",end4-start4,"seconds") pred_cell = [cell_str_to_token_list(i) for i in pred_cell] pred_cell = [re.sub(r'(\d).\s+(\d)', r'\1.\2', i) for i in pred_cell] print(pred_cell) # Combine the table structure and cell content pred_code = build_table_from_html_and_cell(pred_html, pred_cell) pred_code = "".join(pred_code) pred_code = html_table_template(pred_code) # Display the HTML table soup = bs(pred_code) table_code = soup.prettify() print(table_code)