def load_images(input_dir):
import imutil
# Load all images from the input folder
filenames = [os.path.join(input_dir, f) for f in os.listdir(input_dir)]
print('Found {} input files in directory {}'.format(
len(filenames), input_dir))
images = [imutil.load(f, resize_to=(512, 512)) for f in filenames]
return images
def to_array(self, example):
filename = os.path.expanduser(example['filename'])
if not filename.startswith('/'):
filename = os.path.join(DATA_DIR, filename)
box = example.get('box') if self.bounding_box else None
# HACK
#box = (.25, .75, 0, 1)
img = imutil.load(filename)
if self.delete_background:
seg_filename = os.path.expanduser(example['segmentation'])
segmentation = imutil.load(seg_filename)
foreground_mask = np.mean(segmentation, axis=-1) / 255.
img = img * np.expand_dims(foreground_mask, axis=-1)
if self.random_horizontal_flip and random.getrandbits(1):
img = np.flip(img, axis=1)
if self.torch:
img = img.transpose((2, 0, 1))
if self.normalize:
img *= 1.0 / 255
return img
def to_array(self, example):
filename = os.path.expanduser(example['filename'])
if not filename.startswith('/'):
filename = os.path.join(DATA_DIR, filename)
img = imutil.load(filename)
# img = img.transpose((2, 0, 1))
pil_img = Image.fromarray(img.astype('uint8'), 'RGB')
# img *= 1.0 / 255
tensor_img = self.transform(pil_img)
img = tensor_img.numpy()
img *= 1.0 / 255
return img
def filename_to_pixels(self, filename):
# Input is a PNG composed of 6 40x40 monochrome images
# It encodes frames of a game, similar to the SC2 API
# From top-left to bottom-right, maps represent:
# Health, Agent, Small Towers, Big Towers, Friends, Enemies
img = imutil.load(filename, resize_to=None)
assert img.shape == (40*3, 40*2, 3)
# Pytorch convnets require BCHW inputs
channels = np.zeros((6, 40, 40))
channels[0] = img[0:40, 0:40, 0]
channels[1] = img[0:40, 40:80, 0]
channels[2] = img[40:80, 0:40, 0]
channels[3] = img[40:80, 40:80, 0]
channels[4] = img[80:120, 0:40, 0]
channels[5] = img[80:120, 40:80, 0]
# Normalize to [0, 1]
return channels / 255.0
def crop_and_resize(examples):
# resize_name = 'images_x{}'.format(RESIZE)
mkdir(os.path.join(CUB_DIR, 'train'))
mkdir(os.path.join(CUB_DIR, 'val'))
for i, e in enumerate(examples):
filename = e['filename']
img = imutil.load(filename)
# examples[i]['filename'] = filename.replace('images', resize_name)
print(examples[i]['filename'])
# pth, _ = os.path.split(examples[i]['filename'])
# mkdir(pth)
# left, top, box_width, box_height = e['box']
# x0 = int(left)
# x1 = int(left + box_width)
# y0 = int(top)
# y1 = int(top + box_height)
# img = img[y0:y1, x0:x1, :]
# H, W, C = img.shape
# if H >= W:
# height = round(H / W * RESIZE)
# img = imutil.resize(img, resize_width=RESIZE, resize_height=height)
# y0 = (height - RESIZE) // 2
# img = img[y0:y0 + RESIZE, :, :]
# else:
# width = round(W / H * RESIZE)
# img = imutil.resize(img, resize_width=width, resize_height=RESIZE)
# x0 = (width - RESIZE) // 2
# img = img[:, x0:x0 + RESIZE, :]
if e['fold'] == 'train':
filename = filename.replace('images', 'train')
else:
filename = filename.replace('images', 'val')
imutil.show(img, display=False, filename=filename)
return examples
def test_datatype(self):
x = np.random.normal(size=(100, 100, 3))
imutil.show(x, filename='foobar.png')
y = imutil.load('foobar.png', resize_to=(640, 480))
def filename_to_pixels(self, filename):
# Input is a rendered RGB frame of the game
pixels = imutil.load(filename, resize_to=(40,40))
pixels = pixels.transpose(2, 0, 1)
pixels *= 1/255.
return pixels