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