image-classification-service/scripts/client_mock.py

import argparse
import json
from operator import itemgetter

import pdf2image
import requests
from PIL import ImageDraw, ImageFont


def draw_coco_box(draw: ImageDraw.Draw, bbox, klass, proba):
    x1, y1, x2, y2 = itemgetter("x1", "y1", "x2", "y2")(bbox)
    draw.rectangle(((x1, y1), (x2, y2)), outline="red")

    fnt = ImageFont.truetype("Pillow/Tests/fonts/FreeMono.ttf", 30)

    draw.text((x1, y2), text=f"{klass}: {proba:.2f}", fill=(0, 0, 0, 100), font=fnt)


def draw_coco_boxes(image, bboxes, classes, probas):

    draw = ImageDraw.Draw(image)
    for bbox, klass, proba in zip(bboxes, classes, probas):
        draw_coco_box(draw, bbox, klass, proba)

    return image


def annotate(pdf_path, predictions):
    pages = pdf2image.convert_from_path(pdf_path)

    for prd in predictions:
        page_idx, boxes, classes, probas = itemgetter("page_idx", "bboxes", "classes", "probas")(prd)
        page = pages[page_idx]
        image = draw_coco_boxes(page, boxes, classes, probas)
        image.save(f"/tmp/serv_out/{page_idx}.png")


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--pdf_path", required=True)
    args = parser.parse_args()

    return args


def main(args):

    response = requests.post("http://127.0.0.1:5000", data=open(args.pdf_path, "rb"))
    response.raise_for_status()
    predictions = response.json()

    print(json.dumps(predictions, indent=2))
    annotate(args.pdf_path, predictions)


if __name__ == "__main__":
    args = parse_args()
    main(args)