import logging import torch import torch.nn as nn import docling_ibm_models.tableformer.settings as s from docling_ibm_models.tableformer.models.common.base_model import BaseModel from docling_ibm_models.tableformer.models.table04_rs.bbox_decoder_rs import BBoxDecoder from docling_ibm_models.tableformer.models.table04_rs.encoder04_rs import Encoder04 from docling_ibm_models.tableformer.models.table04_rs.transformer_rs import ( Tag_Transformer, ) from docling_ibm_models.tableformer.utils.app_profiler import AggProfiler LOG_LEVEL = logging.WARN # LOG_LEVEL = logging.INFO # LOG_LEVEL = logging.DEBUG class TableModel04_rs(BaseModel, nn.Module): r""" TableNet04Model encoder, dual-decoder model with OTSL+ support """ def __init__(self, config, init_data, device): super(TableModel04_rs, self).__init__(config, init_data, device) self._prof = config["predict"].get("profiling", False) self._device = device # Extract the word_map from the init_data word_map = init_data["word_map"] # Encoder self._enc_image_size = config["model"]["enc_image_size"] self._encoder_dim = config["model"]["hidden_dim"] self._encoder = Encoder04(self._enc_image_size, self._encoder_dim).to(device) tag_vocab_size = len(word_map["word_map_tag"]) td_encode = [] for t in ["ecel", "fcel", "ched", "rhed", "srow"]: if t in word_map["word_map_tag"]: td_encode.append(word_map["word_map_tag"][t]) self._log().debug("td_encode length: {}".format(len(td_encode))) self._log().debug("td_encode: {}".format(td_encode)) self._tag_attention_dim = config["model"]["tag_attention_dim"] self._tag_embed_dim = config["model"]["tag_embed_dim"] self._tag_decoder_dim = config["model"]["tag_decoder_dim"] self._decoder_dim = config["model"]["hidden_dim"] self._dropout = config["model"]["dropout"] self._bbox = config["train"]["bbox"] self._bbox_attention_dim = config["model"]["bbox_attention_dim"] self._bbox_embed_dim = config["model"]["bbox_embed_dim"] self._bbox_decoder_dim = config["model"]["hidden_dim"] self._enc_layers = config["model"]["enc_layers"] self._dec_layers = config["model"]["dec_layers"] self._n_heads = config["model"]["nheads"] self._num_classes = config["model"]["bbox_classes"] self._enc_image_size = config["model"]["enc_image_size"] self._max_pred_len = config["predict"]["max_steps"] self._tag_transformer = Tag_Transformer( device, tag_vocab_size, td_encode, self._decoder_dim, self._enc_layers, self._dec_layers, self._enc_image_size, n_heads=self._n_heads, ).to(device) self._bbox_decoder = BBoxDecoder( device, self._bbox_attention_dim, self._bbox_embed_dim, self._tag_decoder_dim, self._bbox_decoder_dim, self._num_classes, self._encoder_dim, self._dropout, ).to(device) def _log(self): # Setup a custom logger return s.get_custom_logger(self.__class__.__name__, LOG_LEVEL) def mergebboxes(self, bbox1, bbox2): new_w = (bbox2[0] + bbox2[2] / 2) - (bbox1[0] - bbox1[2] / 2) new_h = (bbox2[1] + bbox2[3] / 2) - (bbox1[1] - bbox1[3] / 2) new_left = bbox1[0] - bbox1[2] / 2 new_top = min((bbox2[1] - bbox2[3] / 2), (bbox1[1] - bbox1[3] / 2)) new_cx = new_left + new_w / 2 new_cy = new_top + new_h / 2 bboxm = torch.tensor([new_cx, new_cy, new_w, new_h]) return bboxm def predict(self, imgs, max_steps, k, return_attention=False): r""" Inference. The input image must be preprocessed and transformed. Parameters ---------- img : tensor FloatTensor - torch.Size([1, 3, 448, 448]) Input image for the inference Returns ------- seq : list Predictions for the tags as indices over the word_map outputs_class : tensor(x, 3) Classes of predicted bboxes. x is the number of bboxes. There are 3 bbox classes outputs_coord : tensor(x, 4) Coords of predicted bboxes. x is the number of bboxes. Each bbox is in [cxcywh] format """ AggProfiler().begin("predict_total", self._prof) # Invoke encoder self._tag_transformer.eval() enc_out = self._encoder(imgs) AggProfiler().end("model_encoder", self._prof) word_map = self._init_data["word_map"]["word_map_tag"] n_heads = self._tag_transformer._n_heads # [1, 28, 28, 512] encoder_out = self._tag_transformer._input_filter( enc_out.permute(0, 3, 1, 2) ).permute(0, 2, 3, 1) batch_size = encoder_out.size(0) encoder_dim = encoder_out.size(-1) enc_inputs = encoder_out.view(batch_size, -1, encoder_dim).to(self._device) enc_inputs = enc_inputs.permute(1, 0, 2) positions = enc_inputs.shape[0] encoder_mask = torch.zeros( (batch_size * n_heads, positions, positions), device=self._device ) == torch.ones( (batch_size * n_heads, positions, positions), device=self._device ) # Invoking tag transformer encoder before the loop to save time AggProfiler().begin("model_tag_transformer_encoder", self._prof) encoder_out = self._tag_transformer._encoder(enc_inputs, mask=encoder_mask) AggProfiler().end("model_tag_transformer_encoder", self._prof) decoded_tags = ( torch.LongTensor([word_map[""]]).to(self._device).unsqueeze(1) ) output_tags = [] cache = None tag_H_buf = [] skip_next_tag = True prev_tag_ucel = False line_num = 0 # Populate bboxes_to_merge, indexes of first lcel, and last cell in a span first_lcel = True bboxes_to_merge = {} cur_bbox_ind = -1 bbox_ind = 0 # i = 0 while len(output_tags) < self._max_pred_len: decoded_embedding = self._tag_transformer._embedding(decoded_tags) decoded_embedding = self._tag_transformer._positional_encoding( decoded_embedding ) AggProfiler().begin("model_tag_transformer_decoder", self._prof) decoded, cache = self._tag_transformer._decoder( decoded_embedding, encoder_out, cache, memory_key_padding_mask=encoder_mask, ) AggProfiler().end("model_tag_transformer_decoder", self._prof) # Grab last feature to produce token AggProfiler().begin("model_tag_transformer_fc", self._prof) logits = self._tag_transformer._fc(decoded[-1, :, :]) # 1, vocab_size AggProfiler().end("model_tag_transformer_fc", self._prof) new_tag = logits.argmax(1).item() # STRUCTURE ERROR CORRECTION # Correction for first line xcel... if line_num == 0: if new_tag == word_map["xcel"]: new_tag = word_map["lcel"] # Correction for ucel, lcel sequence... if prev_tag_ucel: if new_tag == word_map["lcel"]: new_tag = word_map["fcel"] # End of generation if new_tag == word_map[""]: output_tags.append(new_tag) decoded_tags = torch.cat( [ decoded_tags, torch.LongTensor([new_tag]).unsqueeze(1).to(self._device), ], dim=0, ) # current_output_len, 1 break output_tags.append(new_tag) # BBOX PREDICTION # MAKE SURE TO SYNC NUMBER OF CELLS WITH NUMBER OF BBOXes if not skip_next_tag: if new_tag in [ word_map["fcel"], word_map["ecel"], word_map["ched"], word_map["rhed"], word_map["srow"], word_map["nl"], word_map["ucel"], ]: # GENERATE BBOX HERE TOO (All other cases)... tag_H_buf.append(decoded[-1, :, :]) if first_lcel is not True: # Mark end index for horizontal cell bbox merge bboxes_to_merge[cur_bbox_ind] = bbox_ind bbox_ind += 1 # Treat horisontal span bboxes... if new_tag != word_map["lcel"]: first_lcel = True else: if first_lcel: # GENERATE BBOX HERE (Beginning of horisontal span)... tag_H_buf.append(decoded[-1, :, :]) first_lcel = False # Mark start index for cell bbox merge cur_bbox_ind = bbox_ind bboxes_to_merge[cur_bbox_ind] = -1 bbox_ind += 1 if new_tag in [word_map["nl"], word_map["ucel"], word_map["xcel"]]: skip_next_tag = True else: skip_next_tag = False # Register ucel in sequence... if new_tag == word_map["ucel"]: prev_tag_ucel = True else: prev_tag_ucel = False decoded_tags = torch.cat( [ decoded_tags, torch.LongTensor([new_tag]).unsqueeze(1).to(self._device), ], dim=0, ) # current_output_len, 1 seq = decoded_tags.squeeze().tolist() if self._bbox: AggProfiler().begin("model_bbox_decoder", self._prof) outputs_class, outputs_coord = self._bbox_decoder.inference( enc_out, tag_H_buf ) AggProfiler().end("model_bbox_decoder", self._prof) else: outputs_class, outputs_coord = None, None outputs_class.to(self._device) outputs_coord.to(self._device) ######################################################################################## # Merge First and Last predicted BBOX for each span, according to bboxes_to_merge ######################################################################################## outputs_class1 = [] outputs_coord1 = [] boxes_to_skip = [] for box_ind in range(len(outputs_coord)): box1 = outputs_coord[box_ind].to(self._device) cls1 = outputs_class[box_ind].to(self._device) if box_ind in bboxes_to_merge: box2 = outputs_coord[bboxes_to_merge[box_ind]].to(self._device) boxes_to_skip.append(bboxes_to_merge[box_ind]) boxm = self.mergebboxes(box1, box2).to(self._device) outputs_coord1.append(boxm) outputs_class1.append(cls1) else: if box_ind not in boxes_to_skip: outputs_coord1.append(box1) outputs_class1.append(cls1) if len(outputs_coord1) > 0: outputs_coord1 = torch.stack(outputs_coord1) else: outputs_coord1 = torch.empty(0) if len(outputs_class1) > 0: outputs_class1 = torch.stack(outputs_class1) else: outputs_class1 = torch.empty(0) outputs_class = outputs_class1 outputs_coord = outputs_coord1 # Do the rest of the steps... AggProfiler().end("predict_total", self._prof) num_tab_cells = seq.count(4) + seq.count(5) num_rows = seq.count(9) self._log().info( "OTSL predicted table cells#: {}; rows#: {}".format(num_tab_cells, num_rows) ) return seq, outputs_class, outputs_coord