def get_prediction_masks(self):
prediction_masks = {}
batch_gen, steps = self.validation_datagen
for batch_id, data in enumerate(batch_gen):
if len(data) != len(self.output_names) + 1:
raise ValueError('incorrect targets provided')
X = data[0]
targets_tensors = data[1:]
if (targets_tensors[0].size()[1] > 1):
targets_tensors = [
target_tensor[:, :1] for target_tensor in targets_tensors
]
if torch.cuda.is_available():
X = Variable(X, volatile=True).cuda()
else:
X = Variable(X, volatile=True)
outputs_batch = self.model(X)
if len(outputs_batch) == len(self.output_names):
for name, output, target in zip(self.output_names,
outputs_batch,
targets_tensors):
if name in self.outputs_to_plot:
prediction = categorize_image(softmax(
output.data.cpu().numpy()),
channel_axis=1)
ground_truth = np.squeeze(target.cpu().numpy(), axis=1)
n_channels = output.data.cpu().numpy().shape[1]
for channel_nr in range(n_channels):
category_id = CATEGORY_IDS[channel_nr]
if category_id != None:
channel_ground_truth = np.where(
ground_truth == channel_nr, 1, 0)
mask_key = '{}_{}'.format(name, category_id)
prediction_masks[mask_key] = np.stack(
[prediction, channel_ground_truth], axis=1)
else:
for name, target in zip(self.output_names, targets_tensors):
if name in self.outputs_to_plot:
prediction = categorize_image(softmax(
outputs_batch.data.cpu().numpy()),
channel_axis=1)
ground_truth = np.squeeze(target.cpu().numpy(), axis=1)
n_channels = outputs_batch.data.cpu().numpy().shape[1]
for channel_nr in range(n_channels):
category_id = CATEGORY_IDS[channel_nr]
if category_id != None:
channel_ground_truth = np.where(
ground_truth == channel_nr, 1, 0)
mask_key = '{}_{}'.format(name, category_id)
prediction_masks[mask_key] = np.stack(
[prediction, channel_ground_truth], axis=1)
break
return prediction_masks
def _transform(self, images):
for image in tqdm(images):
yield categorize_image(image)
def transform(self, images):
categorized_images = []
for image in tqdm(images):
categorized_images.append(categorize_image(image))
return {'categorized_images': categorized_images}
