167 lines
7 KiB
Python
167 lines
7 KiB
Python
from PIL import Image, ImageDraw, ImageColor, ImageFont
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import string
|
|
import random
|
|
import json
|
|
import os
|
|
import math
|
|
|
|
# DEFINE PARAMATERS
|
|
forbidden_letters = set(["I", "O"]) # letters to not use in code charts because can be confused with digits
|
|
px_pt_ratio = 20/29 # according to our image dimensions, 29 point = 20 px
|
|
text_color = ImageColor.getrgb("gray")
|
|
font_type = "arial.ttf"
|
|
tojitter = True # add jitter from a regular grid
|
|
ebuf = 5 # buffer number of pixels to leave from the edges of the image so codecharts are not tangent to image edges
|
|
go_to_image_edges = False # if want to make sure to sample triplets to the very edge of the image (downside: triplets may be more crowded)
|
|
|
|
def point_to_pixel(num):
|
|
return int(num*px_pt_ratio)
|
|
|
|
def pixel_to_point(num):
|
|
return int(num*(1/px_pt_ratio))
|
|
|
|
def generate_rand_letter():
|
|
letter = random.choice(string.ascii_uppercase)
|
|
while letter in forbidden_letters:
|
|
letter = random.choice(string.ascii_uppercase)
|
|
return letter
|
|
|
|
def generate_rand_triplet():
|
|
code = ""
|
|
code += generate_rand_letter()
|
|
# the following code prevents the two digits from being identical or equal to 0
|
|
for i in range(2):
|
|
if i == 0:
|
|
forbidden_num = 0
|
|
else:
|
|
forbidden_num = int(code[i])
|
|
r = list(range(1, forbidden_num)) + list(range(forbidden_num+1, 10))
|
|
code += str(random.choice(r))
|
|
return code
|
|
|
|
def create_codechart(filename,image_width,image_height):
|
|
|
|
font_size = int(image_height*0.0185)
|
|
|
|
# all these parameters depend on font size
|
|
max_triplet_width = font_size*3 # in pixel - max triplet width; used 'W88' as widest triplet code (width~60, height=20)
|
|
max_triplet_height = font_size # the tallest a triplet can be
|
|
d_v = 4*max_triplet_height # vertical distance to maintain b/w triplets in the grid
|
|
d_h = 2*max_triplet_width # horizontal distance to maintain b/w triplets (from start of one triplet to start of another)
|
|
|
|
# make sure that not too much empty space is left over by spacing out triplets
|
|
N_h = int(math.floor((image_width-max_triplet_width-2*ebuf) / float(d_h))) # number of triplets that will be tiled horizontally
|
|
d_h = int(math.floor((image_width-max_triplet_width-2*ebuf) / float(N_h))) # recompute the horizontal dist between triplets to eliminate extra space
|
|
N_v = int(math.floor((image_height-max_triplet_height-2*ebuf) / float(d_v)))
|
|
d_v = int(math.floor((image_height-max_triplet_height-2*ebuf) / float(N_v)))
|
|
# -------------
|
|
|
|
post_jitter_buffer = 6 # small buffer to cover edge case of triplets immediately adjacent to one another (for legibility)
|
|
j_v = int(0.25*(d_v) - post_jitter_buffer/2) # max vertical jitter for one side of a triplet
|
|
j_h = int(0.25*(d_h) - post_jitter_buffer) # max horizontal jitter for on side of a triplet
|
|
|
|
# set up image canvas and font size/style
|
|
img = Image.new('RGB', (image_width, image_height))
|
|
d = ImageDraw.Draw(img)
|
|
try:
|
|
font = ImageFont.truetype(font_type, pixel_to_point(font_size)) # takes in point value
|
|
except OSError:
|
|
print("WARNING: using a different font bc could not find %s on your computer" % font_type)
|
|
font = ImageFont.truetype("/usr/share/fonts/truetype/freefont/FreeSans.ttf", pixel_to_point(font_size))
|
|
|
|
valid_codes = set()
|
|
coordinates = {}
|
|
|
|
# initialize starting locations for triplets on the image grid
|
|
x_init = ebuf
|
|
y_init = ebuf
|
|
|
|
# -------- improvement made after 01/2020 (after TurkEyes paper) --------
|
|
# original problem was grid-like artifacts in the collected data because the grid spacing between consecutive triplets
|
|
# was always similar (despite a small bit of jitter added when using the triplet)
|
|
xoffset = random.choice(list(range(int(d_h/2.0))))
|
|
yoffset = random.choice(list(range(int(d_v/2.0))))
|
|
x_init = x_init+xoffset
|
|
y_init = y_init+yoffset
|
|
# -----------------------------------------------------------------------
|
|
|
|
x = x_init
|
|
while x < image_width-max_triplet_width-ebuf:
|
|
|
|
y = y_init
|
|
while y < image_height-max_triplet_height-ebuf:
|
|
|
|
triplet_code = generate_rand_triplet()
|
|
|
|
# check for if triplet has already been used in image since all codes should be unique
|
|
while triplet_code in valid_codes:
|
|
triplet_code = generate_rand_triplet()
|
|
valid_codes.add(triplet_code)
|
|
|
|
if tojitter:
|
|
# implement jitter to x and y coordinates (note: can turn either of them off)
|
|
min_x = max(ebuf,x-j_h)
|
|
max_x = min(x+j_h+1,image_width-max_triplet_width-ebuf)
|
|
min_y = max(ebuf,y-j_v)
|
|
max_y = min(y+j_v+1,image_height-max_triplet_height-ebuf-2) # a little bit of extra buffer in vertical dimension
|
|
x_range = list(range(min_x, max_x))
|
|
y_range = list(range(min_y, max_y))
|
|
j_x = random.choice(x_range)
|
|
j_y = random.choice(y_range)
|
|
else:
|
|
j_x = x
|
|
j_y = y
|
|
|
|
# writes triplet to image
|
|
d.text((j_x, j_y), triplet_code, text_color, font)
|
|
coordinates[triplet_code] = (j_x, j_y)
|
|
|
|
y_prev = y
|
|
y = y+d_v # regularly sample the image vertically
|
|
|
|
# triplets are not guaranteed to go to edge of image, and gap could be large
|
|
# see if can still squeeze in a triplet without overlapping previous ones (could still be quite close)
|
|
if go_to_image_edges and y >= image_height-max_triplet_height-ebuf:
|
|
y = y_prev + max_triplet_height+j_v+1 + post_jitter_buffer*2
|
|
|
|
x_prev = x
|
|
x = x+d_h # regularly sample the image horizontally
|
|
|
|
if go_to_image_edges and x >= image_width-max_triplet_width-ebuf:
|
|
x = x_prev + max_triplet_width+j_h+1 + post_jitter_buffer*2
|
|
|
|
|
|
img.save(filename)
|
|
return (list(valid_codes), coordinates)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# create some code charts to test this code
|
|
rootdir = './task_data'
|
|
num_codecharts = 3 # generate this many codecharts
|
|
|
|
#image_width = 1920 # in pixel
|
|
#image_height = 1080 # in pixel
|
|
|
|
image_height = 1340 #1344
|
|
image_width = int(1036*image_height/float(1344))
|
|
|
|
# set up directories
|
|
if not os.path.exists(rootdir):
|
|
os.makedirs(rootdir)
|
|
test_dir = os.path.join(rootdir,'TEST')
|
|
if not os.path.exists(test_dir):
|
|
os.makedirs(test_dir)
|
|
|
|
data = {}
|
|
for i in range(num_codecharts):
|
|
filename = os.path.join(test_dir,'CC_%d.jpg'%(i))
|
|
valid_codes, coordinates = create_codechart(filename,image_width,image_height)
|
|
data[filename] = (valid_codes, coordinates)
|
|
|
|
with open(os.path.join(test_dir,'data.json'), 'w') as outfile:
|
|
json.dump(data, outfile)
|
|
|