def process_full_image(image_name, output, multi_run, dataset_folder, class_encodings, post_process):
"""
Helper function to save the output during testing
Parameters
----------
meanIU.avg : float
MeanIU of the model of the evaluated split
multi_run :
image_name: str
name of the image that is saved
output: numpy matrix of size [#C x H x W]
output image at full size
dataset_folder: str
path to the dataset folder
post_process : Boolean
apply post-processing to the output of the network
Returns
-------
mean_iu : float
mean iu of this image
"""
# Load GT
with open(os.path.join(dataset_folder, "test", "gt", image_name), 'rb') as f:
with Image.open(f) as img:
gt = np.array(img)
# Get predictions
if post_process:
# Load original image
with open(os.path.join(dataset_folder, "test", "data", image_name[:-4] + ".JPG"), 'rb') as f:
with Image.open(f) as img:
original_image = np.array(img)
# Apply CRF
# prediction = crf(original_image, output)
# output_encoded = output_to_class_encodings(prediction, class_encodings, perform_argmax=False)
else:
prediction = np.argmax(output, axis=0)
output_encoded = output_to_class_encodings(output, class_encodings)
# Get boundary pixels and adjust the gt_image for the border pixel -> set to background (1)
boundary_mask = gt[:, :, 0].astype(np.uint8) == 128
# Get the ground truth mapping and filter their values for the boundary pixels
target = np.argmax(gt_to_one_hot(gt, class_encodings).numpy(), axis=0)
target[np.logical_and.reduce([boundary_mask, prediction != target, prediction == 0])] = 0 # NOTE: here 0 is 0x1 because it works on the index!
# Compute and record the meanIU of the whole image
_, _, mean_iu, _ = accuracy_segmentation(target, prediction, len(class_encodings))
scalar_label = 'output_{}'.format(image_name) if multi_run is None else 'output_{}_{}'.format(multi_run, image_name)
_save_output_evaluation(class_encodings, output_encoded=output_encoded, gt_image=gt, tag=scalar_label, multi_run=multi_run)
return mean_iu
def train_one_mini_batch(model, criterion, optimizer, input_var,
target_var_argmax, loss_meter, meanIU_meter,
num_classes):
"""
This routing train the model passed as parameter for one mini-batch
Parameters
----------
num_classes:
model : torch.nn.module
The network model being used.
criterion : torch.nn.loss
The loss function used to compute the loss of the model.
optimizer : torch.optim
The optimizer used to perform the weight update.
input_var : torch.autograd.Variable
The input data for the mini-batch
target_var_argmax : torch.autograd.Variable
The target data (labels) for the mini-batch
loss_meter : AverageMeter
Tracker for the overall loss
meanIU_meter : AverageMeter
Tracker for the overall meanIU
Returns
-------
acc : float
Accuracy for this mini-batch
loss : float
Loss for this mini-batch
"""
# Compute output
output = model(input_var)
# Compute and record the loss
loss = criterion(output, target_var_argmax)
try:
loss_meter.update(loss.item(), len(input_var))
except AttributeError:
loss_meter.update(loss.data[0], len(input_var))
output_argmax = np.array(
[np.argmax(o, axis=0) for o in output.data.cpu().numpy()])
target_argmax = target_var_argmax.data.cpu().numpy()
# Compute and record the accuracy
_, _, mean_iu, _ = accuracy_segmentation(target_argmax, output_argmax,
num_classes)
meanIU_meter.update(mean_iu, input_var.size(0))
# Reset gradient
optimizer.zero_grad()
# Compute gradients
loss.backward()
# Perform a step by updating the weights
optimizer.step()
# return acc, loss
return mean_iu, loss
def validate(val_loader, model, criterion, writer, epoch, class_encodings, no_cuda=False, log_interval=10, **kwargs):
"""
The evaluation routine
Parameters
----------
val_loader : torch.utils.data.DataLoader
The dataloader of the evaluation set
model : torch.nn.module
The network model being used
criterion: torch.nn.loss
The loss function used to compute the loss of the model
writer : tensorboardX.writer.SummaryWriter
The tensorboard writer object. Used to log values on file for the tensorboard visualization.
epoch : int
Number of the epoch (for logging purposes)
class_encodings : List
Contains the classes (range of ints)
no_cuda : boolean
Specifies whether the GPU should be used or not. A value of 'True' means the CPU will be used.
log_interval : int
Interval limiting the logging of mini-batches. Default value of 10.
Returns
-------
meanIU.avg : float
MeanIU of the model of the evaluated split
"""
# 'Run' is injected in kwargs at runtime IFF it is a multi-run event
multi_run = kwargs['run'] if 'run' in kwargs else None
num_classes = len(class_encodings)
# Instantiate the counters
batch_time = AverageMeter()
losses = AverageMeter()
meanIU = AverageMeter()
data_time = AverageMeter()
# Switch to evaluate mode (turn off dropout & such )
model.eval()
# Iterate over whole evaluation set
end = time.time()
pbar = tqdm(enumerate(val_loader), total=len(val_loader), unit='batch', ncols=150, leave=False)
for batch_idx, (input, target) in pbar:
# Measure data loading time
data_time.update(time.time() - end)
# Moving data to GPU
if not no_cuda:
input = input.cuda(non_blocking=True)
target = target.cuda(non_blocking=True)
# Compute output
output = model(input)
# Compute and record the loss
loss = criterion(output, target)
losses.update(loss.item(), input.size(0))
# Compute and record the accuracy
_, _, mean_iu_batch, _ = accuracy_segmentation(target.cpu().numpy(), get_argmax(output), num_classes)
meanIU.update(mean_iu_batch, input.size(0))
# Add loss and meanIU to Tensorboard
scalar_label = 'val/mb_loss' if multi_run is None else 'val/mb_loss_{}'.format(multi_run)
writer.add_scalar(scalar_label, loss.item(), epoch * len(val_loader) + batch_idx)
scalar_label = 'val/mb_meanIU' if multi_run is None else 'val/mb_meanIU_{}'.format(multi_run)
writer.add_scalar(scalar_label, mean_iu_batch, epoch * len(val_loader) + batch_idx)
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if batch_idx % log_interval == 0:
pbar.set_description('val epoch [{0}][{1}/{2}]\t'.format(epoch, batch_idx, len(val_loader)))
pbar.set_postfix(Time='{batch_time.avg:.3f}\t'.format(batch_time=batch_time),
Loss='{loss.avg:.4f}\t'.format(loss=losses),
meanIU='{meanIU.avg:.3f}\t'.format(meanIU=meanIU),
Data='{data_time.avg:.3f}\t'.format(data_time=data_time))
# Logging the epoch-wise meanIU
scalar_label = 'val/meanIU' if multi_run is None else 'val/meanIU_{}'.format(multi_run)
writer.add_scalar(scalar_label, meanIU.avg, epoch)
logging.info(_prettyprint_logging_label("val") +
' epoch[{}]: '
'MeanIU={meanIU.avg:.3f}\t'
'Loss={loss.avg:.4f}\t'
'Batch time={batch_time.avg:.3f} ({data_time.avg:.3f} to load data)'
.format(epoch, batch_time=batch_time, data_time=data_time, loss=losses, meanIU=meanIU))
return meanIU.avg
def test(test_loader, model, criterion, writer, epoch, class_encodings, img_names_sizes_dict, dataset_folder,
post_process, no_cuda=False, log_interval=10, **kwargs):
"""
The evaluation routine
Parameters
----------
img_names_sizes_dict: dictionary {str: (int, int)}
Key: gt image name (with extension), Value: image size
test_loader : torch.utils.data.DataLoader
The dataloader of the evaluation set
model : torch.nn.module
The network model being used
criterion: torch.nn.loss
The loss function used to compute the loss of the model
writer : tensorboardX.writer.SummaryWriter
The tensorboard writer object. Used to log values on file for the tensorboard visualization.
epoch : int
Number of the epoch (for logging purposes)
class_encodings : List
Contains the range of encoded classes
img_names_sizes_dict
# TODO
dataset_folder : str
# TODO
post_process : Boolean
apply post-processing to the output of the network
no_cuda : boolean
Specifies whether the GPU should be used or not. A value of 'True' means the CPU will be used.
log_interval : int
Interval limiting the logging of mini-batches. Default value of 10.
Returns
-------
meanIU.avg : float
MeanIU of the model of the evaluated split
"""
# 'Run' is injected in kwargs at runtime IFF it is a multi-run event
multi_run = kwargs['run'] if 'run' in kwargs else None
num_classes = len(class_encodings)
# Instantiate the counters
batch_time = AverageMeter()
losses = AverageMeter()
meanIU = AverageMeter()
data_time = AverageMeter()
# Switch to evaluate mode (turn off dropout & such )
model.eval()
# Iterate over whole evaluation set
end = time.time()
# Need to store the images currently being processes
canvas = {}
pbar = tqdm(enumerate(test_loader), total=len(test_loader), unit='batch', ncols=150, leave=False)
for batch_idx, (input, target) in pbar:
# Unpack input
input, top_left_coordinates, test_img_names = input
# Measure data loading time
data_time.update(time.time() - end)
# Moving data to GPU
if not no_cuda:
input = input.cuda(non_blocking=True)
target = target.cuda(non_blocking=True)
# Compute output
output = model(input)
# Compute and record the loss
loss = criterion(output, target)
losses.update(loss.item(), input.size(0))
# Compute and record the batch meanIU
_, _, mean_iu_batch, _ = accuracy_segmentation(target.cpu().numpy(), get_argmax(output), num_classes)
# Add loss and meanIU to Tensorboard
scalar_label = 'test/mb_loss' if multi_run is None else 'test/mb_loss_{}'.format(multi_run)
writer.add_scalar(scalar_label, loss.item(), epoch * len(test_loader) + batch_idx)
scalar_label = 'test/mb_meanIU' if multi_run is None else 'test/mb_meanIU_{}'.format(multi_run)
writer.add_scalar(scalar_label, mean_iu_batch, epoch * len(test_loader) + batch_idx)
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if batch_idx % log_interval == 0:
pbar.set_description('test epoch [{0}][{1}/{2}]\t'.format(epoch, batch_idx, len(test_loader)))
pbar.set_postfix(Time='{batch_time.avg:.3f}\t'.format(batch_time=batch_time),
Loss='{loss.avg:.4f}\t'.format(loss=losses),
meanIU='{meanIU.avg:.3f}\t'.format(meanIU=meanIU),
Data='{data_time.avg:.3f}\t'.format(data_time=data_time))
# Output needs to be patched together to form the complete output of the full image
# patches are returned as a sliding window over the full image, overlapping sections are averaged
for patch, x, y, img_name in zip(output.data.cpu().numpy(), top_left_coordinates[0].numpy(), top_left_coordinates[1].numpy(), test_img_names):
# Is a new image?
if not img_name in canvas:
# Create a new image of the right size filled with NaNs
canvas[img_name] = np.empty((num_classes, *img_names_sizes_dict[img_name]))
canvas[img_name].fill(np.nan)
# Add the patch to the image
canvas[img_name] = merge_patches(patch, (x, y), canvas[img_name])
# Save the image when done
if not np.isnan(np.sum(canvas[img_name])):
# Save the final image
mean_iu = process_full_image(img_name, canvas[img_name], multi_run, dataset_folder, class_encodings, post_process)
# Update the meanIU
meanIU.update(mean_iu, 1)
# Remove the entry
canvas.pop(img_name)
logging.info("\nProcessed image {} with mean IU={}".format(img_name, mean_iu))
# Canvas MUST be empty or something was wrong with coverage of all images
assert len(canvas) == 0
# Logging the epoch-wise meanIU
scalar_label = 'test/mb_meanIU' if multi_run is None else 'test/mb_meanIU_{}'.format(multi_run)
writer.add_scalar(scalar_label, meanIU.avg, epoch)
logging.info(_prettyprint_logging_label("test") +
' epoch[{}]: '
'MeanIU={meanIU.avg:.3f}\t'
'Loss={loss.avg:.4f}\t'
'Batch time={batch_time.avg:.3f} ({data_time.avg:.3f} to load data)'
.format(epoch, batch_time=batch_time, data_time=data_time, loss=losses, meanIU=meanIU))
return meanIU.avg
def process_full_image(image_name, output, multi_run, dataset_folder,
class_encodings, post_process):
"""
Helper function to save the output during testing
Parameters
----------
meanIU.avg : float
MeanIU of the model of the evaluated split
multi_run :
image_name: str
name of the image that is saved
output: numpy matrix of size [#C x H x W]
output image at full size
dataset_folder: str
path to the dataset folder
post_process : Boolean
apply post-processing to the output of the network
Returns
-------
mean_iu : float
mean iu of this image
"""
# Load GT
with open(os.path.join(dataset_folder, "test", "gt", image_name),
'rb') as f:
with Image.open(f) as img:
gt = np.array(img)
# Get predictions
if post_process:
# Load original image
with open(
os.path.join(dataset_folder, "test", "data",
image_name[:-4] + ".JPG"), 'rb') as f:
with Image.open(f) as img:
original_image = np.array(img)
# # Apply CRF
# prediction = crf(original_image, output)
# output_encoded = output_to_class_encodings(prediction, class_encodings, perform_argmax=False)
else:
prediction = np.argmax(output, axis=0)
output_encoded = output_to_class_encodings(output, class_encodings)
# Get the ground truth mapping
target = np.argmax(gt_to_one_hot(gt, class_encodings).numpy(), axis=0)
# Compute and record the meanIU of the whole image
_, _, mean_iu, _ = accuracy_segmentation(target, prediction,
len(class_encodings))
scalar_label = 'output_{}'.format(
image_name) if multi_run is None else 'output_{}_{}'.format(
multi_run, image_name)
_save_output_evaluation(class_encodings,
output_encoded=output_encoded,
tag=scalar_label,
multi_run=multi_run)
return mean_iu