ExCeipt / Layoutlmv3_inference /annotate_image.py
Scezui's picture
1.supported multiple file uploads
144dfb5
raw
history blame
2.12 kB
import os
from PIL import Image, ImageDraw, ImageFont
from .utils import image_label_2_color
def get_flattened_output(docs):
print("Running Flattened Output")
flattened_output = []
annotation_key = 'output'
for doc in docs:
flattened_output_item = {annotation_key: []}
doc_annotation = doc[annotation_key]
for i, span in enumerate(doc_annotation):
if len(span['words']) > 1:
for span_chunk in span['words']:
flattened_output_item[annotation_key].append(
{
'label': span['label'],
'text': span_chunk['text'],
'words': [span_chunk]
}
)
else:
flattened_output_item[annotation_key].append(span)
flattened_output.append(flattened_output_item)
return flattened_output
def annotate_image(image_path, annotation_object):
print("Annotating Images")
img = None
image = Image.open(image_path).convert('RGBA')
tmp = image.copy()
label2color = image_label_2_color(annotation_object)
overlay = Image.new('RGBA', tmp.size, (0, 0, 0)+(0,))
draw = ImageDraw.Draw(overlay)
font = ImageFont.load_default()
predictions = [span['label'] for span in annotation_object['output']]
boxes = [span['words'][0]['box'] for span in annotation_object['output']]
for prediction, box in zip(predictions, boxes):
draw.rectangle(box, outline=label2color[prediction],
width=3, fill=label2color[prediction]+(int(255*0.33),))
draw.text((box[0] + 10, box[1] - 10), text=prediction,
fill=label2color[prediction], font=font)
img = Image.alpha_composite(tmp, overlay)
img = img.convert("RGB")
image_name = os.path.basename(image_path)
image_name = image_name[:image_name.find('.')]
output_folder = 'static/temp/img_display/'
os.makedirs(output_folder, exist_ok=True)
img.save(os.path.join(output_folder, f'{image_name}.jpg_inference.jpg'))