59 lines
1.5 KiB
Python
59 lines
1.5 KiB
Python
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)
|