visrecall/WebInterface/Front-end/generate-experiment-files/generate_sentinels.py

from PIL import Image, ImageDraw, ImageColor, ImageFont
import matplotlib.pyplot as plt
import numpy as np
import string
import random
import json 
import generate_codecharts as gc
import os

# FIXED parameters
text_color = ImageColor.getrgb("white")
font_type = "arial.ttf"
px_pt_ratio = 20/29 # according to our image dimensions, 29 point = 20 px
valid_target_types=["red_dot", "fix_cross", "img"]

def make_sentinel(codechart_filename,sentinel_filename,image_width,image_height,border_padding,target_type="red_dot", target_im_dir=""):
    # border_padding used to guarantee that chosen sentinel location is not too close to border to be hard to spot
    
    font_size = int(image_height*0.0278)
    correct_codes = []
    
    if target_type not in valid_target_types: 
        raise RuntimeError("target_type must be one of %s" % valid_target_types.__str__())
    valid_codes, coordinates = gc.create_codechart(codechart_filename,image_width,image_height)
    # pick random code 
    r = list(range(0, len(valid_codes)))
    index = random.choice(r)
    triplet = valid_codes[index]
    triplet_coordinate = coordinates[triplet]
    # to make sure that the cross is visible
    while (triplet_coordinate[0] <= border_padding or triplet_coordinate[0] >= image_width-border_padding) \
    or (triplet_coordinate[1] <= border_padding or triplet_coordinate[1] >=image_height-border_padding): 
        index = random.choice(r)
        triplet = valid_codes[index]
        triplet_coordinate = coordinates[triplet]
    # check bg color
    if target_type == "fix_cross": 
        bg_color = 126
    else: 
        bg_color = 255
    # create and save cross sentinel image
    img = Image.new('RGB', (image_width, image_height), (bg_color, bg_color, bg_color))
    d = ImageDraw.Draw(img)
    try: 
        font = ImageFont.truetype(font_type, gc.pixel_to_point(font_size)) # takes in point value
    except OSError:
        print("WARNING: using different font bc could not find %s" % font_type)
        font = ImageFont.truetype("/usr/share/fonts/truetype/freefont/FreeSans.ttf", gc.pixel_to_point(font_size))

    if target_type == "fix_cross":
        plot_coord = (triplet_coordinate[0]+font_size, triplet_coordinate[1]) # offset cross location to the center of the triplet
        d.text(plot_coord, '+', text_color, font)
    elif target_type == "red_dot":
        d.ellipse((triplet_coordinate[0], triplet_coordinate[1], triplet_coordinate[0]+font_size*2, triplet_coordinate[1]+font_size*2), fill = 'red', outline ='red')
    elif target_type == "img": 
        if not target_im_dir:
            raise RuntimeError("No im dir provided for sentinel targets")
        # Get a list of images in the target im dir
        images = os.listdir(target_im_dir)
        target = Image.open(os.path.join(target_im_dir, random.choice(images)))
        # resize the target 
        width = 200
        height = int(target.height*width/target.width)
        target = target.resize((width, height))
        plot_coord = (triplet_coordinate[0]-int(width/2), triplet_coordinate[1]-int(height/2))
        img.paste(target, plot_coord)
        
        # correct_codes lie within the sentinel width
        for ii in range(len(valid_codes)):
            dist = np.linalg.norm(np.array(coordinates[valid_codes[ii]])-np.array(triplet_coordinate))
            if dist <= width/2.0:
                correct_codes.append(valid_codes[ii]);
                
        pass
    else: 
        raise RuntimeError("target_type %s does not exist" % target_type)
    img.save(sentinel_filename)
    D = {'correct_code':triplet, 'coordinate':triplet_coordinate, 'correct_codes':correct_codes}
    D_full = {'correct_code':triplet, 'coordinate':triplet_coordinate, \
              'valid_codes':valid_codes, 'coordinates':coordinates, 'codechart_file':codechart_filename,\
              'correct_codes':correct_codes}
    return D,D_full

def generate_sentinels(sentinel_image_dir,sentinel_CC_dir,num_buckets,start_bucket_at,sentinel_images_per_bucket,\
                       image_width,image_height,border_padding,target_type,target_im_dir=""):
    
    # Set up directories
    if not os.path.exists(sentinel_image_dir):
        os.makedirs(sentinel_image_dir)
    if not os.path.exists(sentinel_CC_dir):
        os.makedirs(sentinel_CC_dir)
        
    # Start generating sentinels
    img_num_offset = (start_bucket_at-1)*sentinel_images_per_bucket # start at a new index id
    for b in range(num_buckets):

        image_bucket_dir = os.path.join(sentinel_image_dir,'bucket%d'%(start_bucket_at+b))
        if not os.path.exists(image_bucket_dir):
            os.makedirs(image_bucket_dir)

        CC_bucket_dir = os.path.join(sentinel_CC_dir,'bucket%d'%(start_bucket_at+b))
        if not os.path.exists(CC_bucket_dir):
            os.makedirs(CC_bucket_dir)

        data = {} # save to a json the filename, the coordinate of the + cross, and the triplet at that coordinate
        data_with_coords = {} # also save a list of other valid triplets and coordinates (for analysis)

        print('Populating %s with %d sentinel images'%(image_bucket_dir,sentinel_images_per_bucket))
        print('Populating %s with %d corresponding codecharts'%(CC_bucket_dir,sentinel_images_per_bucket))
        for i in range(sentinel_images_per_bucket):
            img_num = img_num_offset + b*sentinel_images_per_bucket + i + 1
            # generate random code chart
            codechart_filename = os.path.join(CC_bucket_dir,'sentinel_CC_%d.jpg'%(img_num))
            sentinel_filename = os.path.join(image_bucket_dir,'sentinel_image_%d.jpg'%(img_num))

            D,D_full = make_sentinel(codechart_filename,sentinel_filename,image_width,image_height,border_padding,target_type, target_im_dir)

            data[sentinel_filename] = D
            data_with_coords[sentinel_filename] = D_full

        with open(os.path.join(image_bucket_dir,'sentinel_codes.json'), 'w') as outfile: 
            json.dump(data, outfile)
        print('Writing out %s'%(os.path.join(image_bucket_dir,'sentinel_codes.json')))

        with open(os.path.join(image_bucket_dir,'sentinel_codes_full.json'), 'w') as outfile: 
            json.dump(data_with_coords, outfile)
        print('Writing out %s'%(os.path.join(image_bucket_dir,'sentinel_codes_full.json')))


if __name__ == "__main__":
    
    # Set these parameters
    sentinel_images_per_bucket = 500 
    num_buckets = 1 
    start_bucket_at = 0 # to avoid overwriting the existing buckets
    
    image_width = 1920 
    image_height = 1080 
    border_padding = 100 # don't put fixation cross in this region of the image
    rootdir = './task_data'

    target_type = "img" 
    target_im_dir = "sentinel_target_images"

    sentinel_image_dir = os.path.join(rootdir,'sentinel_images')
    sentinel_CC_dir = os.path.join(rootdir,'sentinel_CC')
    
    generate_sentinels(sentinel_image_dir,sentinel_CC_dir,num_buckets,start_bucket_at,sentinel_images_per_bucket,\
                       image_width,image_height,border_padding,target_type, target_im_dir)