def load_batch(self, index):
# print index
shuffled = lambda seq, rnd=random.random: sorted(seq, key=lambda _: rnd())
should_horizontally_flip = random.getrandbits(1)
near_distance = self.nearfarsplit[index]
far_distance = self.far_max
far_distance = np.max([near_distance + 1000, far_distance])
selection = [0] + shuffled(range(1, near_distance))[0:5] + shuffled(range(near_distance, far_distance))[0:105]
spatial_transform = Compose([Scale(112),
CenterCrop(112),
ToTensor(),
Normalize(get_mean(), [1, 1, 1])])
near_indices = self.nearest_indices[index]
image_crops_list = []
prepped_tensors = []
for idx, near_idx in enumerate(near_indices[selection]):
image_name = self.image_root + self.dataset[near_idx][0] + self.dataset[near_idx][1].split('_')[1] + '.png'
raw_imgs = scipy.misc.imread(image_name)[np.newaxis,:]
r1, g1, b1 = [115, 108, 99]
r2, g2, b2, = get_mean()
red, green, blue = raw_imgs[:, :, :, 0], raw_imgs[:, :, :, 1], raw_imgs[:, :, :, 2]
mask = (red == r1) & (green == g1) & (blue == b1)
raw_imgs[:, :, :, :3][mask] = [r2, g2, b2]
image_crops = self.augment(raw_imgs, do_flip=should_horizontally_flip)
image_crops_list += [image_crops]
prepped_images = [spatial_transform(Image.fromarray(image_crop)) for image_crop in image_crops][0]
prepped_tensor = prepped_images
prepped_tensors += [prepped_tensor]
# labels = ['anchor', 'similar', 'similar', 'similar', 'similar', 'similar', 'different', 'different', 'different']
# image_crops_list = image_crops_list[0:9]
# for frame in [0]:
# fig = plt.figure()
# for axes_idx in range(1,10):
# fig.add_subplot(3, 3, axes_idx)
# plt.imshow(image_crops_list[axes_idx-1][frame])
# plt.title(labels[axes_idx-1])
# plt.axis('off')
# plt.show()
batch = torch.stack(prepped_tensors, 0)
batch2 = torch.zeros(batch.size())
batch2[:, 0] = batch[:, 2]
batch2[:, 1] = batch[:, 1]
batch2[:, 2] = batch[:, 0]
return batch2
def load_batch(self, index):
# print index
selection = [index]
spatial_transform = Compose([
Scale(112),
CenterCrop(112),
ToTensor(),
Normalize(get_mean(), [1, 1, 1])
])
image_crops_list = []
prepped_tensors = []
for idx, near_idx in enumerate(selection):
image_name = self.image_root + self.dataset[near_idx][
0] + self.dataset[near_idx][1].split('_')[1] + '.png'
raw_imgs = scipy.misc.imread(image_name)[np.newaxis, :]
r1, g1, b1 = [115, 108, 99]
r2, g2, b2, = get_mean()
red, green, blue = raw_imgs[:, :, :,
0], raw_imgs[:, :, :,
1], raw_imgs[:, :, :, 2]
mask = (red == r1) & (green == g1) & (blue == b1)
raw_imgs[:, :, :, :3][mask] = [r2, g2, b2]
image_crops = self.augment(raw_imgs)
image_crops_list += [image_crops]
prepped_images = [
spatial_transform(Image.fromarray(image_crop))
for image_crop in image_crops
][0]
prepped_tensor = prepped_images
prepped_tensors += [prepped_tensor]
batch = torch.stack(prepped_tensors, 0)
batch2 = torch.zeros(batch.size())
batch2[:, 0] = batch[:, 2]
batch2[:, 1] = batch[:, 1]
batch2[:, 2] = batch[:, 0]
return batch2