Layoutlmv3_inference/inference_handler.py

from .utils import load_model,load_processor,normalize_box,compare_boxes,adjacent
from .annotate_image import get_flattened_output,annotate_image
from PIL import Image,ImageDraw, ImageFont
import logging
import torch
import json
import os


logger = logging.getLogger(__name__)

class ModelHandler(object):
    def __init__(self):
        self.model = None
        self.model_dir = None
        self.device = 'cpu'
        self.error = None
        self.initialized = False
        self._raw_input_data = None
        self._processed_data = None
        self._images_size = None
        self.counter = 0


    def initialize(self, context):
        try:
            logger.info("Loading transformer model")
            self._context = context
            properties = self._context
            self.model_dir = properties.get("model_dir")
            self.model = self.load(self.model_dir)
            self.initialized = True
        except Exception as e:
            logger.error(f"Error initializing model: {str(e)}")
            self.error = str(e)

    def preprocess(self, batch):
        try:
            inference_dict = batch
            self._raw_input_data = inference_dict
            processor = load_processor()
            images = [Image.open(path).convert("RGB")
                      for path in inference_dict['image_path']]
            self._images_size = [img.size for img in images]
            words = inference_dict['words']
            boxes = [[normalize_box(box, images[i].size[0], images[i].size[1])
                      for box in doc] for i, doc in enumerate(inference_dict['bboxes'])]
            encoded_inputs = processor(
                images, words, boxes=boxes, return_tensors="pt", padding="max_length", truncation=True)
            self._processed_data = encoded_inputs
            return encoded_inputs
        except Exception as e:
            logger.error(f"Error in preprocessing: {str(e)}")
            self.error = str(e)
            return None

    def load(self, model_dir):
        try:
            model = load_model(model_dir)
            return model
        except Exception as e:
            logger.error(f"Error loading LayoutLMv3 model: {str(e)}")
            self.error = str(e)
            return None

    def inference(self, model_input):
        try:
            with torch.no_grad():
                inference_outputs = self.model(**model_input)
                predictions = inference_outputs.logits.argmax(-1).tolist()
            results = []
            for i in range(len(predictions)):
                tmp = dict()
                tmp[f'output_{i}'] = predictions[i]
                results.append(tmp)
            return [results]
        except Exception as e:
            logger.error(f"Error in inference: {str(e)}")
            self.error = str(e)
            return None

    def postprocess(self, inference_output):
        try:
            docs = []
            k = 0
            for page, doc_words in enumerate(self._raw_input_data['words']):
                doc_list = []
                width, height = self._images_size[page]
                for i, doc_word in enumerate(doc_words, start=0):
                    word_tagging = None
                    word_labels = []
                    word = dict()
                    word['id'] = k
                    k += 1
                    word['text'] = doc_word
                    word['pageNum'] = page + 1
                    word['box'] = self._raw_input_data['bboxes'][page][i]
                    _normalized_box = normalize_box(
                        self._raw_input_data['bboxes'][page][i], width, height)
                    for j, box in enumerate(self._processed_data['bbox'].tolist()[page]):
                        if compare_boxes(box, _normalized_box):
                            if self.model.config.id2label[inference_output[0][page][f'output_{page}'][j]] != 'O':
                                word_labels.append(
                                    self.model.config.id2label[inference_output[0][page][f'output_{page}'][j]][2:])
                            else:
                                word_labels.append('other')
                    if word_labels != []:
                        word_tagging = word_labels[0] if word_labels[0] != 'other' else word_labels[-1]
                    else:
                        word_tagging = 'other'
                    word['label'] = word_tagging
                    word['pageSize'] = {'width': width, 'height': height}
                    if word['label'] != 'other':
                        doc_list.append(word)
                spans = []
                def adjacents(entity): return [
                    adj for adj in doc_list if adjacent(entity, adj)]
                output_test_tmp = doc_list[:]
                for entity in doc_list:
                    if adjacents(entity) == []:
                        spans.append([entity])
                        output_test_tmp.remove(entity)

                while output_test_tmp != []:
                    span = [output_test_tmp[0]]
                    output_test_tmp = output_test_tmp[1:]
                    while output_test_tmp != [] and adjacent(span[-1], output_test_tmp[0]):
                        span.append(output_test_tmp[0])
                        output_test_tmp.remove(output_test_tmp[0])
                    spans.append(span)

                output_spans = []
                for span in spans:
                    if span[0]['label'] in ['ITEMS', 'PRICE']:
                        output_span = {"text": ' '.join([entity['text'] for entity in span]) + '|',
                                    "label": span[0]['label'],
                                    "words": [{
                                        'id': entity['id'],
                                        'box': entity['box'],
                                        'text': entity['text']
                                    } for entity in span]
                                    }
                    else:
                        output_span = {"text": ' '.join([entity['text'] for entity in span]),
                                    "label": span[0]['label'],
                                    "words": [{
                                        'id': entity['id'],
                                        'box': entity['box'],
                                        'text': entity['text']
                                    } for entity in span]
                                    }
                    output_spans.append(output_span)
                docs.append({f'output': output_spans})
            return [json.dumps(docs, ensure_ascii=False)]

        except Exception as e:
            logger.error(f"Error in postprocessing: {str(e)}")
            self.error = str(e)
            return None


    def handle(self, data, context):
        """
        Call preprocess, inference, and post-process functions
        :param data: input data
        :param context: mms context
        """
        try:
            if not self.initialized:
                self.initialize(context)

            if data is None:
                return None

            model_input = self.preprocess(data)
            if model_input is None:
                return None

            model_out = self.inference(model_input)
            if model_out is None:
                return None

            inference_out = self.postprocess(model_out)[0]

            # # Write the original inference output to a JSON file
            # with open('temp/LayoutlMV3InferenceOutput.json', 'w') as inf_out:
            #     inf_out.write(inference_out)

            # Load the original inference output from the JSON file
            inference_out_list = json.loads(inference_out)
            flattened_output_list = get_flattened_output(inference_out_list)
            print('Ready for Annotation')
            for i, flattened_output in enumerate(flattened_output_list):
                annotate_image(data['image_path'][i], flattened_output)

            # Create the labeled directory if it doesn't exist
            labeled_dir = 'static/temp/labeled'
            os.makedirs(labeled_dir, exist_ok=True)

            # Loop through the list and create a new file for each JSON object
            for index, output_json in enumerate(inference_out_list):
                # Generate a filename based on the index
                output_filename = os.path.join(labeled_dir, f'Output_{index}.json')
                # Write the JSON object to the file
                with open(output_filename, 'w') as output_file:
                    json.dump(output_json, output_file)
        
        except Exception as e:
            logger.error(f"Error handling request: {str(e)}")
            self.error = str(e)

                        
_service = ModelHandler()


def handle(data, context):
    if not _service.initialized:
        _service.initialize(context)

    if data is None:
        return None

    return _service.handle(data, context)