#!/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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment