def _prediction_to_label_tiff_image(prediction):
"""Save the prediction to a label image (TIFF)"""
# Convert to label image
label_img = one_hot_stack_to_label_image(
prediction,
first_index_is_background=True,
channels_first=True,
dtype=np.uint16)
return label_img
def get_or_load_mask_at_index(self, index: int):
"""Get mask for requested index from cache or load it.
In case of failure, the last_error_message property contains
the error message.
@param index: Index of the mask to get or load.
@return the mask if it could be loaded properly, None otherwise.
"""
mask = None
if index in self._masks:
mask = self._masks[index]
else:
# Load it
if index < len(self._mask_names):
mask_file_name = self._mask_names[index]
if mask_file_name is not None:
# What is the file type?
extension = str(Path(mask_file_name).suffix).lower()
if extension == ".tif" or extension == ".tiff":
# Read the file
try:
mask = imread(mask_file_name)
except Exception as e:
# Set the error message
self._error_message = str(e)
# Report failure
return None
# Set the mask type
self._mask_type = MaskType.TIFF_LABELS
# Set the number of classes if not set
if self._num_classes == 0:
self._num_classes = len(np.unique(mask))
elif extension == ".npy":
# Read the file
try:
mask = np.load(mask_file_name)
except Exception as e:
# Set the error message
self._error_message = str(e)
# Report failure
return None
# Set the mask type
if mask.ndim == 2:
self._mask_type = MaskType.NUMPY_LABELS
else:
self._mask_type = MaskType.NUMPY_ONE_HOT
# Set the number of classes if not set
if self._num_classes == 0:
if self._mask_type == MaskType.NUMPY_LABELS:
self._num_classes = len(np.unique(mask))
else:
self._num_classes = mask.shape[0]
# Convert one-hot stack to label image for display
if self._mask_type == MaskType.NUMPY_ONE_HOT:
# Binarize the dataset
mask[mask > 0.5] = True
# Change into a labels image (for display)
mask = one_hot_stack_to_label_image(
mask,
first_index_is_background=True,
channels_first=True,
dtype=np.int32)
elif extension == ".h5":
# Read the file
mask_file = h5py.File(mask_file_name, 'r')
# Get the dataset
datasets = list(mask_file.keys())
if len(datasets) != 1:
# Set the error message
self._error_message = f"Unexpected number of datasets in file {mask_file_name}"
# Report failure
return None
# Dataset name
dataset = datasets[0]
# Get the dataset
mask = np.array(mask_file[dataset])
# Close the HDF5 file
mask_file.close()
if mask.ndim == 2:
# Set the error message
self._error_message = f"The dataset in file {mask_file_name} must have "\
f"dimensions (num_classes x height x width."
# Report failure
return None
# Set the type
self._mask_type = MaskType.H5_ONE_HOT
# Set the number of classes if not set
if self._num_classes == 0:
self._num_classes = mask.shape[0]
# Binarize the dataset
mask[mask > 0.5] = True
# Change into a labels image (for display)
mask = one_hot_stack_to_label_image(
mask,
first_index_is_background=True,
channels_first=True,
dtype=np.int32)
else:
self._mask_type = MaskType.UNKNOWN
# Set the error message
self._error_message = f"Unexpected mask file {mask_file_name}"
# Report failure
return None
# Make sure the mask is an int32
mask = mask.astype(np.int32)
# Add it to the cache
self._masks[index] = mask
return mask