Skip to content

Instantly share code, notes, and snippets.

@anishathalye
Created December 10, 2019 18:01
Show Gist options
  • Save anishathalye/77ee771c208f888c7fe098dca32e8a22 to your computer and use it in GitHub Desktop.
Save anishathalye/77ee771c208f888c7fe098dca32e8a22 to your computer and use it in GitHub Desktop.
#!/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