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)