from layoutlm_preprocess import *
from PIL import Image, ImageDraw, ImageFont

def iob_to_label(label):
  #print(label)  
  if label != 'O':
    print(label[2:])  
    return label[2:]
  else:
    return ""

image = Image.open('/datadrive/unilm/layoutlm/deprecated/examples/seq_labeling/images/542ff60a-invoice131.jpg')
image = image.convert("RGB")
#image.show()
image_path='/datadrive/unilm/layoutlm/deprecated/examples/seq_labeling/images/542ff60a-invoice131.jpg'


print(image_path)
labels = get_labels("/datadrive/unilm/layoutlm/deprecated/examples/seq_labeling/invoice_data/labels.txt")
num_labels = len(labels)
image, words, boxes, actual_boxes = preprocess(image_path)
#print(words)
#print(actual_boxes)
print(len(words))
print(len(actual_boxes))
model_path='/datadrive/unilm/layoutlm/deprecated/examples/seq_labeling/output/layoutlm_invoice.pt'
model=model_load(model_path,num_labels)
#print(model)
word_level_predictions, final_boxes=convert_to_features(image, words, boxes, actual_boxes, model)

#print(word_level_predictions)

#print(final_boxes)

print(len(word_level_predictions)) #291
print(len(final_boxes)) #291

draw = ImageDraw.Draw(image)
font = ImageFont.load_default()

label2color = {'invoice_due_date':'yellow','invoice_date':'orange','total_amount':'blue','telephone_number':'green','invoice_number':'red','address':'blue','':'violet'}

for prediction, box in zip(word_level_predictions, final_boxes):
    predicted_label = iob_to_label(label_map[prediction]).lower()
    print(predicted_label)
    draw.rectangle(box, outline=label2color[predicted_label])
    draw.text((box[0] + 10, box[1] - 10), text=predicted_label, fill=label2color[predicted_label], font=font)

#image.show()   
image.save(f'/datadrive/unilm/layoutlm/deprecated/examples/seq_labeling/output_test/infer_img/542ff60a-invoice131.jpg', "JPEG")