screenshot-service/Docker/annotate_header.py

87 lines
3.0 KiB
Python

# test with:
# docker image build -t headless-fox Docker && docker run -v $PWD/static:/opt/static -ti headless-fox time python annotate_header.py screenshot.png "/opt/static/output.png" "content-type" "Tutaj jest content-type" "etag" "Tutaj jest etag z długim opisem co ma wiele słów i wychodzi poza network inspector"
import os
import sys
import pytesseract
import uuid
import json
from pytesseract import Output
from PIL import Image, ImageDraw, ImageFont
output_file_relative = sys.argv[
1
] # this is also the existing source screenshot to annotate. It will be updated in-place
output_file = "/opt/static/" + output_file_relative
domain = sys.argv[2]
needles = sys.argv[3:]
base_url = os.getenv("BASE_URL")
# generator
def partition(lst, size):
for i in range(0, len(lst), size):
yield lst[i : i + size]
# print(d)
with Image.open(output_file) as im:
x_offset = 2054
y_offset = 313
cropped = im.crop((x_offset, y_offset, 2875, 1558))
cropped_filename = "/opt/static/" + uuid.uuid4().hex + ".png"
cropped.save(cropped_filename)
d = pytesseract.image_to_data(cropped_filename, output_type=Output.DICT)
os.remove(cropped_filename)
draw = ImageDraw.Draw(im)
n_boxes = len(d["level"])
font = ImageFont.truetype("/usr/share/fonts/noto/NotoSansDisplay-Medium.ttf", 48)
found_needles = []
for i in range(n_boxes):
(x, y, w, h, text) = (
d["left"][i],
d["top"][i],
d["width"][i],
d["height"][i],
d["text"][i],
)
for [needle, comment] in partition(needles, 2):
if needle.lower() in text.lower():
found_needles.append(needle)
# modify y so it's aligned not with the top of the text, but with the midline
y = y + h / 2
radius = 30
# offset both y and x
y = y + y_offset
x = x + x_offset
fill = "red"
line_length = 200
draw.regular_polygon(
((x - radius - 5, y), radius), n_sides=3, rotation=270, fill=fill
)
draw.line((x - radius - 5, y, x - line_length, y), fill=fill, width=10)
text_w, text_h = draw.textsize(comment, font)
text_padding = 10
draw.rectangle(
[
(x - line_length - text_w - text_padding, y - text_h / 2),
(x - line_length + text_padding, y + text_h / 2),
],
fill="white",
)
draw.text(
(x - line_length - 10, y),
comment,
fill=fill,
anchor="rm",
font=font,
)
im.save(output_file, "PNG")
print(json.dumps({"new_file":
{"url": base_url + "/static/" + output_file_relative,
"domain": domain,
"found_headers": found_needles}}))