2022-05-29 22:00:28 +02:00
# test with:
2022-06-15 20:27:15 +02:00
# 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"
2022-05-29 22:00:28 +02:00
import os
import sys
import pytesseract
import uuid
from pytesseract import Output
2022-06-15 20:27:15 +02:00
from PIL import Image , ImageDraw , ImageFont
2022-05-29 22:00:28 +02:00
2022-06-15 20:27:15 +02:00
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
needles = sys . argv [ 2 : ]
base_url = os . getenv ( " BASE_URL " )
# generator
def partition ( lst , size ) :
for i in range ( 0 , len ( lst ) , size ) :
yield lst [ i : i + size ]
2022-05-29 22:00:28 +02:00
# print(d)
2022-06-15 20:27:15 +02:00
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 " ] )
print ( needles )
font = ImageFont . truetype ( " /usr/share/fonts/noto/NotoSansDisplay-Medium.ttf " , 48 )
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 ( ) :
print ( " needle " , needle , text , x , y )
# 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 " )
2022-06-17 09:37:59 +02:00
print ( ' { " new_file " : " ' + base_url + " /static/ " + output_file_relative + " \" } " )