File size: 1,941 Bytes
6df3c38 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
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 = 'inferenced/'
os.makedirs(output_folder, exist_ok=True)
img.save(os.path.join(output_folder, f'{image_name}_inference.jpg'))
|