init

.gitignore

@@ -0,0 +1,3 @@
+__pycache__
+**test**
+storage

app.py

@@ -0,0 +1,35 @@
+import streamlit as st
+from PIL import Image
+
+#Trick to not init function multitime
+if "model" not in st.session_state:
+    print("INIT MODEL")
+    from src.model import Model
+    st.session_state.model = Model()
+    print("DONE INIT MODEL")
+
+st.set_page_config(page_title="VQA", layout="wide")
+hide_menu_style = """
+<style>
+footer {visibility: hidden;}
+</style>
+"""
+st.markdown(hide_menu_style, unsafe_allow_html= True)
+
+
+image = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png", "webp", ])
+
+if image:
+    bytes_data = image.getvalue()
+    with open("test.png", "wb") as f:
+        f.write(bytes_data)
+    f.close()
+    st.session_state.image = "test.png"
+
+if 'image' in st.session_state:
+    st.image(st.session_state.image)
+    question = st.text_input("Question: ")
+    if question:
+        answer = st.session_state.model.inference(st.session_state.image, question)
+        st.write(f"Answer: {answer}")
+        

requirements.txt

@@ -0,0 +1,5 @@
+paddlepaddle>=2.3.1 
+paddleocr==2.6.1.3 
+vietocr>=0.3.8
+pillow==9.5.0
+torchvision==0.18.0

src/feature_extraction.py

@@ -0,0 +1,193 @@
+
+import torch
+import requests
+from PIL import Image, ImageFont, ImageDraw, ImageTransform
+from transformers import AutoImageProcessor, ViTModel, AutoTokenizer, T5EncoderModel
+from utils.config import Config
+from src.ocr import OCRDetector
+
+
+class ViT:
+    def __init__(self) -> None:
+        self.processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+        self.model = ViTModel.from_pretrained("google/vit-base-patch16-224-in21k")
+        self.model.to(Config.device)
+
+    def extraction(self, image_url):
+        if image_url.startswith("https://"):
+            images = Image.open(requests.get(image_url, stream=True).raw).convert("RGB")
+        else:
+            images = Image.open(image_url).convert("RGB")
+
+        inputs = self.processor(images, return_tensors="pt").to(Config.device)
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        last_hidden_states = outputs.last_hidden_state
+        attention_mask = torch.ones((last_hidden_states.shape[0], last_hidden_states.shape[1]))
+
+        return last_hidden_states.to(Config.device), attention_mask.to(Config.device)
+
+    def pooling_extraction(self, image):
+        image_inputs = self.processor(image, return_tensors="pt").to(Config.device)
+
+        with torch.no_grad():
+            image_outputs = self.model(**image_inputs)
+            image_pooler_output = image_outputs.pooler_output
+            image_pooler_output = torch.unsqueeze(image_pooler_output, 0)
+            image_attention_mask = torch.ones((image_pooler_output.shape[0], image_pooler_output.shape[1]))
+
+        return image_pooler_output.to(Config.device), image_attention_mask.to(Config.device)
+
+class OCR:
+    def __init__(self) -> None:
+        self.ocr_detector = OCRDetector()
+
+    def extraction(self, image_dir):
+
+        ocr_results = self.ocr_detector.text_detector(image_dir)
+        if not ocr_results:
+            print("NOT OCR1")
+
+            return "", [], []
+
+        ocrs = self.post_process(ocr_results)
+
+        if not ocrs:
+
+            return "", [], []
+
+        ocrs.reverse()
+
+        boxes = []
+        texts = []
+        for idx, ocr in enumerate(ocrs):
+            boxes.append(ocr["box"])
+            texts.append(ocr["text"])
+
+        groups_box, groups_text, paragraph_boxes = OCR.group_boxes(boxes, texts)
+        for temp in groups_text:
+            print("OCR: ", temp)
+
+        texts = [" ".join(group_text) for group_text in groups_text]
+        ocr_content = "<extra_id_0>".join(texts)
+        ocr_content = ocr_content.lower()
+        ocr_content = " ".join(ocr_content.split())
+        ocr_content = "<extra_id_0>" + ocr_content
+
+
+        return ocr_content, groups_box, paragraph_boxes
+
+    def post_process(self,ocr_results):
+        ocrs = []
+        for result in ocr_results:
+            text = result["text"]
+            # if len(text) <=2:
+            #   continue
+            # if len(set(text.replace(" ", ""))) <=2:
+            #   continue
+            box = result["box"]
+
+            # (x1, y1), (x2, y2), (x3, y3), (x4, y4) = box
+            # w = x2 - x1
+            # h = y4 - y1
+            # if h > w:
+            #   continue
+
+            # if w*h < 300:
+            #   continue
+
+            ocrs.append(
+                {"text": text.lower(),
+                "box": box}
+            )
+        return ocrs
+
+    @staticmethod
+    def cut_image_polygon(image, box):
+        (x1, y1), (x2, y2), (x3, y3), (x4, y4) = box
+        w = x2 - x1
+        h = y4 - y1
+        scl = h//7
+        new_box = [max(x1-scl,0), max(y1 - scl, 0)], [x2+scl, y2-scl], [x3+scl, y3+scl], [x4-scl, y4+scl]
+        (x1, y1), (x2, y2), (x3, y3), (x4, y4) = new_box
+        # Define 8-tuple with x,y coordinates of top-left, bottom-left, bottom-right and top-right corners and apply
+        transform = [x1, y1, x4, y4, x3, y3, x2, y2]
+        result = image.transform((w,h), ImageTransform.QuadTransform(transform))
+        return result
+
+
+    @staticmethod
+    def check_point_in_rectangle(box, point, padding_devide):
+      (x1, y1), (x2, y2), (x3, y3), (x4, y4) = box
+      x_min = min(x1, x4)
+      x_max = max(x2, x3)
+
+      padding = (x_max-x_min)//padding_devide
+      x_min = x_min - padding
+      x_max = x_max + padding
+
+      y_min = min(y1, y2)
+      y_max = max(y3, y4)
+
+      y_min = y_min - padding
+      y_max = y_max + padding
+
+      x, y = point
+
+      if x >= x_min and x <= x_max and y >= y_min and y <= y_max:
+        return True
+
+      return False
+
+    @staticmethod
+    def check_rectangle_overlap(rec1, rec2, padding_devide):
+      for point in rec1:
+        if OCR.check_point_in_rectangle(rec2, point, padding_devide):
+          return True
+
+      for point in rec2:
+        if OCR.check_point_in_rectangle(rec1, point, padding_devide):
+          return True
+
+      return False
+
+    @staticmethod
+    def group_boxes(boxes, texts):
+      groups = []
+      groups_text = []
+      paragraph_boxes = []
+      processed = []
+      boxes_cp = boxes.copy()
+      for i, (box, text) in enumerate(zip(boxes_cp, texts)):
+        (x1, y1), (x2, y2), (x3, y3), (x4, y4) = box
+
+        if i not in processed:
+          processed.append(i)
+        else:
+          continue
+
+        groups.append([box])
+        groups_text.append([text])
+        for j, (box2, text2) in enumerate(zip(boxes_cp[i+1:], texts[i+1:])):
+          if j+i+1 in processed:
+            continue
+          padding_devide = len(groups[-1])*4
+          is_overlap = OCR.check_rectangle_overlap(box, box2, padding_devide)
+          if is_overlap:
+            (xx1, yy1), (xx2, yy2), (xx3, yy3), (xx4, yy4) = box2
+            processed.append(j+i+1)
+            groups[-1].append(box2)
+            groups_text[-1].append(text2)
+            new_x1 = min(x1, xx1)
+            new_y1 = min(y1, yy1)
+            new_x2 = max(x2, xx2)
+            new_y2 = min(y2, yy2)
+            new_x3 = max(x3, xx3)
+            new_y3 = max(y3, yy3)
+            new_x4 = min(x4, xx4)
+            new_y4 = max(y4, yy4)
+
+            box = [(new_x1, new_y1), (new_x2, new_y2), (new_x3, new_y3), (new_x4, new_y4)]
+
+        paragraph_boxes.append(box)
+      return groups, groups_text, paragraph_boxes

src/model.py

@@ -0,0 +1,445 @@
+import transformers
+from transformers.models.t5.modeling_t5 import *
+from transformers.models.t5.modeling_t5 import T5Stack
+import os
+import gdown
+import torch
+from typing import *
+from transformers import T5ForConditionalGeneration, AutoTokenizer
+from utils.config import Config
+from src.feature_extraction import ViT, OCR
+
+_CONFIG_FOR_DOC = "T5Config"
+_CHECKPOINT_FOR_DOC = "google-t5/t5-small"
+
+class CustomT5Stack(T5Stack):
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        inputs_embeds=None,
+        head_mask=None,
+        cross_attn_head_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        images_embeds=None,
+    ):
+        # Model parallel
+        if self.model_parallel:
+            torch.cuda.set_device(self.first_device)
+            self.embed_tokens = self.embed_tokens.to(self.first_device)
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            err_msg_prefix = "decoder_" if self.is_decoder else ""
+            raise ValueError(
+                f"You cannot specify both {err_msg_prefix}input_ids and {err_msg_prefix}inputs_embeds at the same time"
+            )
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            input_ids = input_ids.view(-1, input_shape[-1])
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            err_msg_prefix = "decoder_" if self.is_decoder else ""
+            raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds")
+
+        if inputs_embeds is None:
+            if self.embed_tokens is None:
+                raise ValueError("You have to initialize the model with valid token embeddings")
+            inputs_embeds = self.embed_tokens(input_ids)
+            if not self.is_decoder and images_embeds is not None:
+                inputs_embeds = torch.concat([inputs_embeds, images_embeds], dim=1)
+                input_shape = inputs_embeds.size()[:-1]
+
+        batch_size, seq_length = input_shape
+
+        # required mask seq length can be calculated via length of past
+        mask_seq_length = past_key_values[0][0].shape[2] + seq_length if past_key_values is not None else seq_length
+
+        if use_cache is True:
+            if not self.is_decoder:
+                raise ValueError(f"`use_cache` can only be set to `True` if {self} is used as a decoder")
+
+        # initialize past_key_values with `None` if past does not exist
+        if past_key_values is None:
+            past_key_values = [None] * len(self.block)
+
+        if attention_mask is None:
+            attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
+
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape)
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if self.is_decoder and encoder_hidden_states is not None:
+            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(
+                    encoder_hidden_shape, device=inputs_embeds.device, dtype=torch.long
+                )
+            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_extended_attention_mask = None
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        # Prepare head mask if needed
+        head_mask = self.get_head_mask(head_mask, self.config.num_layers)
+        cross_attn_head_mask = self.get_head_mask(cross_attn_head_mask, self.config.num_layers)
+        present_key_value_states = () if use_cache else None
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        all_cross_attentions = () if (output_attentions and self.is_decoder) else None
+        position_bias = None
+        encoder_decoder_position_bias = None
+
+        hidden_states = self.dropout(inputs_embeds)
+
+        for i, (layer_module, past_key_value) in enumerate(zip(self.block, past_key_values)):
+            layer_head_mask = head_mask[i]
+            cross_attn_layer_head_mask = cross_attn_head_mask[i]
+            # Model parallel
+            if self.model_parallel:
+                torch.cuda.set_device(hidden_states.device)
+                # Ensure that attention_mask is always on the same device as hidden_states
+                if attention_mask is not None:
+                    attention_mask = attention_mask.to(hidden_states.device)
+                if position_bias is not None:
+                    position_bias = position_bias.to(hidden_states.device)
+                if encoder_hidden_states is not None:
+                    encoder_hidden_states = encoder_hidden_states.to(hidden_states.device)
+                if encoder_extended_attention_mask is not None:
+                    encoder_extended_attention_mask = encoder_extended_attention_mask.to(hidden_states.device)
+                if encoder_decoder_position_bias is not None:
+                    encoder_decoder_position_bias = encoder_decoder_position_bias.to(hidden_states.device)
+                if layer_head_mask is not None:
+                    layer_head_mask = layer_head_mask.to(hidden_states.device)
+                if cross_attn_layer_head_mask is not None:
+                    cross_attn_layer_head_mask = cross_attn_layer_head_mask.to(hidden_states.device)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    layer_module.forward,
+                    hidden_states,
+                    extended_attention_mask,
+                    position_bias,
+                    encoder_hidden_states,
+                    encoder_extended_attention_mask,
+                    encoder_decoder_position_bias,
+                    layer_head_mask,
+                    cross_attn_layer_head_mask,
+                    None,  # past_key_value is always None with gradient checkpointing
+                    use_cache,
+                    output_attentions,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask=extended_attention_mask,
+                    position_bias=position_bias,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_extended_attention_mask,
+                    encoder_decoder_position_bias=encoder_decoder_position_bias,
+                    layer_head_mask=layer_head_mask,
+                    cross_attn_layer_head_mask=cross_attn_layer_head_mask,
+                    past_key_value=past_key_value,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                )
+
+            # layer_outputs is a tuple with:
+            # hidden-states, key-value-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights)
+            if use_cache is False:
+                layer_outputs = layer_outputs[:1] + (None,) + layer_outputs[1:]
+
+            hidden_states, present_key_value_state = layer_outputs[:2]
+
+            # We share the position biases between the layers - the first layer store them
+            # layer_outputs = hidden-states, key-value-states (self-attention position bias), (self-attention weights),
+            # (cross-attention position bias), (cross-attention weights)
+            position_bias = layer_outputs[2]
+            if self.is_decoder and encoder_hidden_states is not None:
+                encoder_decoder_position_bias = layer_outputs[4 if output_attentions else 3]
+            # append next layer key value states
+            if use_cache:
+                present_key_value_states = present_key_value_states + (present_key_value_state,)
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[3],)
+                if self.is_decoder:
+                    all_cross_attentions = all_cross_attentions + (layer_outputs[5],)
+
+            # Model Parallel: If it's the last layer for that device, put things on the next device
+            if self.model_parallel:
+                for k, v in self.device_map.items():
+                    if i == v[-1] and "cuda:" + str(k) != self.last_device:
+                        hidden_states = hidden_states.to("cuda:" + str(k + 1))
+
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+
+        # Add last layer
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    present_key_value_states,
+                    all_hidden_states,
+                    all_attentions,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=present_key_value_states,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+class CustomT5ForConditionalGeneration(T5ForConditionalGeneration):
+    @add_start_docstrings_to_model_forward(T5_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        decoder_attention_mask: Optional[torch.BoolTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        decoder_head_mask: Optional[torch.FloatTensor] = None,
+        cross_attn_head_mask: Optional[torch.Tensor] = None,
+        encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        images_embeds: Optional[torch.FloatTensor] = None,
+    ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ...,
+            config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for
+            labels in `[0, ..., config.vocab_size]`
+
+        Returns:
+
+        Examples:
+
+        ```python
+        >>> from transformers import AutoTokenizer, T5ForConditionalGeneration
+
+        >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small")
+        >>> model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-small")
+
+        >>> # training
+        >>> input_ids = tokenizer("The <extra_id_0> walks in <extra_id_1> park", return_tensors="pt").input_ids
+        >>> labels = tokenizer("<extra_id_0> cute dog <extra_id_1> the <extra_id_2>", return_tensors="pt").input_ids
+        >>> outputs = model(input_ids=input_ids, labels=labels)
+        >>> loss = outputs.loss
+        >>> logits = outputs.logits
+
+        >>> # inference
+        >>> input_ids = tokenizer(
+        ...     "summarize: studies have shown that owning a dog is good for you", return_tensors="pt"
+        ... ).input_ids  # Batch size 1
+        >>> outputs = model.generate(input_ids)
+        >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+        >>> # studies have shown that owning a dog is good for you.
+        ```"""
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask
+        if head_mask is not None and decoder_head_mask is None:
+            if self.config.num_layers == self.config.num_decoder_layers:
+                warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning)
+                decoder_head_mask = head_mask
+
+        # Encode if needed (training, first prediction pass)
+        if encoder_outputs is None:
+            # Convert encoder inputs in embeddings if needed
+            encoder_outputs = self.encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                head_mask=head_mask,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                images_embeds=images_embeds
+            )
+        elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
+            encoder_outputs = BaseModelOutput(
+                last_hidden_state=encoder_outputs[0],
+                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
+            )
+
+        hidden_states = encoder_outputs[0]
+
+        if self.model_parallel:
+            torch.cuda.set_device(self.decoder.first_device)
+
+        if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:
+            # get decoder inputs from shifting lm labels to the right
+            decoder_input_ids = self._shift_right(labels)
+
+        # Set device for model parallelism
+        if self.model_parallel:
+            torch.cuda.set_device(self.decoder.first_device)
+            hidden_states = hidden_states.to(self.decoder.first_device)
+            if decoder_input_ids is not None:
+                decoder_input_ids = decoder_input_ids.to(self.decoder.first_device)
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(self.decoder.first_device)
+            if decoder_attention_mask is not None:
+                decoder_attention_mask = decoder_attention_mask.to(self.decoder.first_device)
+
+        # Decode
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            inputs_embeds=decoder_inputs_embeds,
+            past_key_values=past_key_values,
+            encoder_hidden_states=hidden_states,
+            encoder_attention_mask=attention_mask,
+            head_mask=decoder_head_mask,
+            cross_attn_head_mask=cross_attn_head_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = decoder_outputs[0]
+
+        # Set device for model parallelism
+        if self.model_parallel:
+            torch.cuda.set_device(self.encoder.first_device)
+            self.lm_head = self.lm_head.to(self.encoder.first_device)
+            sequence_output = sequence_output.to(self.lm_head.weight.device)
+
+        if self.config.tie_word_embeddings:
+            # Rescale output before projecting on vocab
+            # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586
+            sequence_output = sequence_output * (self.model_dim**-0.5)
+
+        lm_logits = self.lm_head(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss(ignore_index=-100)
+            # move labels to correct device to enable PP
+            labels = labels.to(lm_logits.device)
+            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
+            # TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666
+
+        if not return_dict:
+            output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs
+            return ((loss,) + output) if loss is not None else output
+
+        return Seq2SeqLMOutput(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+        )
+        
+transformers.models.t5.modeling_t5.T5Stack = CustomT5Stack
+transformers.models.t5.modeling_t5.T5ForConditionalGeneration = CustomT5ForConditionalGeneration
+transformers.T5ForConditionalGeneration = CustomT5ForConditionalGeneration
+
+
+class Model:
+    def __init__(self) -> None:
+        os.makedirs("storage", exist_ok=True)
+        
+        if not os.path.exists("storage/vlsp_transfomer_vietocr.pth"):
+            print("DOWNLOADING model")
+            gdown.download(Config.model_url, output="storage/vlsp_transfomer_vietocr.pth")
+        self.vit5_tokenizer = AutoTokenizer.from_pretrained("VietAI/vit5-base")
+        self.model = T5ForConditionalGeneration.from_pretrained("truong-xuan-linh/VQA-vit5",
+                                                                revision=Config.revision,
+                                                                output_attentions=True)
+        self.model.to(Config.device)
+
+        self.vit = ViT()
+        self.ocr = OCR()
+
+    def get_inputs(self, image_dir: str, question: str):
+        #VIT
+        image_feature, image_mask = self.vit.extraction(image_dir)
+
+        ocr_content, groups_box, paragraph_boxes = self.ocr.extraction(image_dir)
+        print("Input: ", question + " " + ocr_content)
+        #VIT5
+        input_ = self.vit5_tokenizer(question + " " + ocr_content,
+                    padding="max_length",
+                    truncation=True,
+                    max_length=Config.question_maxlen + Config.ocr_maxlen,
+                    return_tensors="pt")
+
+        input_ids = input_.input_ids
+        attention_mask = input_.attention_mask
+        mask = torch.cat((attention_mask, image_mask), 1)
+        return {
+                  "input_ids": input_ids,
+                  "attention_mask": mask,
+                  "images_embeds": image_feature,
+              }
+
+    def inference(self, image_dir: str, question: str):
+        inputs = self.get_inputs(image_dir, question)
+        with torch.no_grad():
+            input_ids = inputs["input_ids"]
+            attention_mask = inputs["attention_mask"]
+            images_embeds = inputs["images_embeds"]
+            generated_ids = self.model.generate(
+                                input_ids=input_ids, \
+                                attention_mask=attention_mask, \
+                                images_embeds=images_embeds, \
+                                num_beams=2,
+                                max_length=Config.answer_maxlen
+                            )
+
+            pred_answer = self.vit5_tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+        return pred_answer

src/ocr.py

@@ -0,0 +1,79 @@
+from paddleocr import PaddleOCR
+from vietocr.tool.config import Cfg
+from vietocr.tool.predictor import Predictor
+from utils.config import Config
+import requests
+import numpy as np
+from PIL import Image, ImageTransform
+
+class OCRDetector:
+  def __init__(self) -> None:
+    self.paddle_ocr = PaddleOCR(lang='en',
+                                use_angle_cls=False,
+                                use_gpu=True if Config.device == "cpu" else False,
+                                show_log=False )
+    # config['weights'] = './weights/transformerocr.pth'
+
+    vietocr_config = Cfg.load_config_from_name('vgg_transformer')
+    vietocr_config['weights'] = Config.ocr_path
+    vietocr_config['cnn']['pretrained']=False
+    vietocr_config['device'] = Config.device
+    vietocr_config['predictor']['beamsearch']=False
+    self.viet_ocr = Predictor(vietocr_config)
+
+  def find_box(self, image):
+    '''Xác định box dựa vào mô hình paddle_ocr'''
+    result = self.paddle_ocr.ocr(image, cls = False, rec=False)
+    result = result[0]
+    # Extracting detected components
+    boxes = result #[res[0] for res in result]
+    boxes = np.array(boxes).astype(int)
+
+    # scores = [res[1][1] for res in result]
+    return boxes
+
+  def cut_image_polygon(self, image, box):
+    (x1, y1), (x2, y2), (x3, y3), (x4, y4) = box
+    w = x2 - x1
+    h = y4 - y1
+    scl = h//7
+    new_box = [max(x1-scl,0), max(y1 - scl, 0)], [x2+scl, y2-scl], [x3+scl, y3+scl], [x4-scl, y4+scl]
+    (x1, y1), (x2, y2), (x3, y3), (x4, y4) = new_box
+    # Define 8-tuple with x,y coordinates of top-left, bottom-left, bottom-right and top-right corners and apply
+    transform = [x1, y1, x4, y4, x3, y3, x2, y2]
+    result = image.transform((w,h), ImageTransform.QuadTransform(transform))
+    return result
+
+  def vietnamese_text(self, boxes, image):
+    '''Xác định text dựa vào mô hình viet_ocr'''
+    results = []
+    for box in boxes:
+      try:
+        cut_image = self.cut_image_polygon(image, box)
+        # cut_image = Image.fromarray(np.uint8(cut_image))
+        text, score = self.viet_ocr.predict(cut_image, return_prob=True)
+        if score > Config.vietocr_threshold:
+          results.append({"text": text,
+                        "score": score,
+                        "box": box})
+      except:
+        continue
+    return results
+
+  #Merge
+  def text_detector(self, image_path):
+    if image_path.startswith("https://"):
+        image = Image.open(requests.get(image_path, stream=True).raw).convert("RGB")
+    else:
+        image = Image.open(image_path).convert("RGB")
+    # np_image = np.array(image)
+
+    boxes = self.find_box(image_path)
+    if not boxes.any():
+        return None
+
+    results = self.vietnamese_text(boxes, image)
+    if results != []:
+        return results
+    else:
+        return None

utils/config.py

@@ -0,0 +1,13 @@
+class Config:
+    device = "cpu"
+    model_url = "https://drive.google.com/uc?id=1WqyPOxgTj9vdnEQl_TJwr_U_hdQeI1iz"
+    ocr_path = "storage/vlsp_transfomer_vietocr.pth"
+    model_path = "storage/vit_base_vit5_base_v2_1.3197_0.4732_3.5212.pt"
+    question_maxlen = 32
+    vietocr_threshold = 0.5
+    answer_maxlen = 56
+    ocr_maxlen = 128
+    ocr_maxobj = 10000
+    num_ocr = 32
+    num_beams = 3
+    revision = "version_2_with_extra_id_0"