image-classification-service/scripts/client_mock.py

import argparse
import json
from operator import itemgetter

import pdf2image
import requests
from PIL import ImageDraw


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


def draw_coco_boxes(image, bboxes, classes):

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

    return image


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

    for prd in predictions:
        page_idx, boxes, classes = itemgetter("page_idx", "bboxes", "classes")(prd)
        page = pages[page_idx]
        image = draw_coco_boxes(page, boxes, classes)
        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://0.0.0.0:8080", 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)