tile-highlight-images.py

#!/usr/bin/env python

import numpy as np
import PIL.Image
import scipy.misc
import sys
import os
import random

def color(hex_repr):
    n = int(hex_repr, 16)
    r = (n >> 16) / 255.0
    g = (n >> 8 & 255) / 255.0
    b = (n & 255) / 255.0
    ret = [r, g, b]
    return ret


CORRECT_COLOR = color('27ae60')
ADV_COLOR = color('d32f2f')
MIS_COLOR = color('000000')


def main():
    data_dir, output_filename, width, height, tile_dim, border_dim, border_line_width, extra_seed = sys.argv[1:]
    width = int(width)
    height = int(height)
    tile_dim = int(tile_dim)
    border_dim = int(border_dim)
    border_line_width = int(border_line_width)
    total = width * height
    final_tile_size = (tile_dim + 2*border_dim)

    output_dir = os.path.dirname(output_filename)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # height, width, depth
    output_image = np.ones((height*final_tile_size, width*final_tile_size, 3))

    # choose pseudorandom subset
    files = [os.path.splitext(os.path.basename(i))[0] for i in os.listdir(data_dir) if i.endswith('.txt')]
    # this should be reproducible!!
    random.seed(int('labsix', 36) + 10000*height + 100*width + int(extra_seed))
    random.shuffle(files)
    files = files[:height*width]

    index = 0
    for ih in range(height):
        for iw in range(width):
            image = get_image(os.path.join(data_dir, '%s.png' % files[index]), tile_dim)
            with open(os.path.join(data_dir, '%s.txt' % files[index])) as f:
                classification = f.read().strip()
            ph = ih * final_tile_size
            pw = iw * final_tile_size
            output_image[ph:ph+final_tile_size, pw:pw+final_tile_size, :] = \
                    format_image(image, classification, tile_dim, border_dim, border_line_width, final_tile_size)
            index += 1

    scipy.misc.imsave(output_filename, output_image)

def get_image(path, tile_dim):
    img = PIL.Image.open(path)
    img = img.resize((tile_dim, tile_dim), resample=PIL.Image.BILINEAR)
    return np.asarray(img).astype(np.float32)/255.0

def format_image(image, classification, tile_dim, border_dim, border_line_width, final_tile_size):
    out = np.ones((final_tile_size, final_tile_size, 3))
    start = border_dim - border_line_width
    end = final_tile_size - start
    if classification == 'correct':
        out[start:end, start:end, :] = CORRECT_COLOR
    elif classification == 'adversarial':
        out[start:end, start:end, :] = ADV_COLOR
    else:
        assert classification == 'misclassified'
        out[start:end, start:end, :] = MIS_COLOR
    out[border_dim:border_dim+tile_dim, border_dim:border_dim+tile_dim, :] = image
    return out

if __name__ == '__main__':
    main()