diff --git a/src/napari_convpaint/convpaint_model.py b/src/napari_convpaint/convpaint_model.py index 6c13f94..cb58a4b 100644 --- a/src/napari_convpaint/convpaint_model.py +++ b/src/napari_convpaint/convpaint_model.py @@ -871,7 +871,47 @@ def predict_probas(self, image, in_channels=None, skip_norm=False, use_dask=Fals probas = self._predict(image, add_seg=False, in_channels=in_channels, skip_norm=skip_norm, use_dask=use_dask, fe_use_device=fe_use_device) return probas - + + def get_instances(self, image, in_channels=None, skip_norm=False, use_dask=False, fe_use_device=None, + min_size=0, classes=None, per_plane=False, warn=True): + """ + Creates instance masks from the semantic segmentation of an image. + + Parameters + ---------- + image : np.ndarray or list[np.ndarray] + Image to create instances for or list of images + in_channels : list[int], optional + List of channels to use for instance creation + skip_norm : bool, optional + Whether to skip normalization of the images before instance creation + use_dask : bool, optional + Whether to use dask for parallel processing + fe_use_device : str, optional + Device policy for feature extractor ("auto", "gpu", "cpu") + min_size : int, optional + Minimum size of instances to keep (default: 0) + classes : list[int], optional + List of classes to create instances for (default: all classes) + warn : bool, optional + Whether to show warnings (default: True) + + Returns + ---------- + instances : np.ndarray or list[np.ndarray] + Instance masks created from the semantic segmentation or list of instance masks for each image if input is a list. + """ + _, seg = self._predict(image, add_seg=True, in_channels=in_channels, + skip_norm=skip_norm, use_dask=use_dask, fe_use_device=fe_use_device) + + return utils.create_instances_from_semantic( + segmentations=seg, + min_size=min_size, + classes=classes, + per_plane=per_plane, + warn=warn + ) + def get_feature_image(self, data, in_channels=None, skip_norm=False, pca_components=0, kmeans_clusters=0, diff --git a/src/napari_convpaint/convpaint_widget.py b/src/napari_convpaint/convpaint_widget.py index 649f069..f12911a 100644 --- a/src/napari_convpaint/convpaint_widget.py +++ b/src/napari_convpaint/convpaint_widget.py @@ -383,7 +383,7 @@ def __init__(self, napari_viewer, parent=None, third_party=False): # Reset to initial state self.reset_class_names_btn = QPushButton('Reset to default') self.classes_layout.addWidget(self.reset_class_names_btn, len(self.initial_names)+3, 0, 1, 10) - self.btn_class_distribution_annot = QPushButton('Show class distribution (in annotation)') + self.btn_class_distribution_annot = QPushButton('Show class distribution (in annotations layer)') self.classes_layout.addWidget(self.btn_class_distribution_annot, len(self.initial_names)+4, 0, 1, 10) # Create the class names @@ -523,6 +523,26 @@ def __init__(self, napari_viewer, parent=None, third_party=False): self.check_add_probas.setChecked(self.add_probas) self.advanced_output_group.glayout.addWidget(self.check_add_probas, 0, 1, 1, 1) + # Checkbox for adding instances + self.check_add_instances = QCheckBox('Instances') + self.check_add_instances.setChecked(self.add_instances) + self.advanced_output_group.glayout.addWidget(self.check_add_instances, 0, 2, 1, 1) + + # Make output checkbox columns expand with widget width. + self.advanced_output_group.glayout.setColumnStretch(0, 1) + self.advanced_output_group.glayout.setColumnStretch(1, 1) + self.advanced_output_group.glayout.setColumnStretch(2, 1) + + # Instance size option + self.inst_min_size_label = QLabel('Instance min_num_pix (0 = ignore)') + self.text_inst_min_size = QtWidgets.QLineEdit() + self.text_inst_min_size.setStyleSheet("font-size: 12px;") + self.text_inst_min_size.setPlaceholderText('e.g. 100') + self.text_inst_min_size.setText(self.inst_min_size) + self.advanced_output_group.glayout.addWidget(self.inst_min_size_label, 1, 0, 1, 2) + self.advanced_output_group.glayout.addWidget(self.text_inst_min_size, 1, 2, 1, 1) + # self.advanced_output_group.glayout.setColumnStretch(1, 2) + # Button to add features for the current plane self.btn_add_features = QPushButton('Get features image') self.advanced_unsupervised_group.glayout.addWidget(self.btn_add_features, 2, 0, 1, 2) @@ -754,6 +774,9 @@ def _setup_init_tooltips(self): self.btn_switch_axes.setToolTip('Switch first two axes of a 4D input image (to match the convention to have channels first).') self.check_add_seg.setToolTip('Add a layer with the predicted segmentation as output (= highest class probability).') self.check_add_probas.setToolTip('Add a layer with class probabilities as output.') + self.check_add_instances.setToolTip('Add a layer with instance masks as output.') + for w in [self.inst_min_size_label, self.text_inst_min_size]: + w.setToolTip('Minimum number of pixels an instance must have to be kept. Set to 0 to disable filtering and watershedding (uses connected components instead).') self.btn_add_features.setToolTip('Add a layer with the features extracted for the current plane.') self.btn_add_features_stack.setToolTip('Add a layer with the features extracted for the whole stack.') for w in [self.pca_label, self.text_features_pca]: @@ -823,7 +846,7 @@ def _remove_init_tooltips(self): self.check_auto_select_annot, # self.text_annot_prefix, self.btn_train_on_selected, self.radio_img_training, self.radio_global_training, self.radio_single_training, # self.check_cont_training, self.btn_class_distribution_trained, self.btn_reset_training, self.check_use_dask, self.channels_label, - self.text_input_channels, self.btn_switch_axes, self.check_add_seg, self.check_add_probas, self.btn_add_features, self.btn_add_features_stack, + self.text_input_channels, self.btn_switch_axes, self.check_add_seg, self.check_add_probas, self.check_add_instances, self.inst_min_size_label, self.text_inst_min_size, self.btn_add_features, self.btn_add_features_stack, self.pca_label, self.text_features_pca, self.kmeans_label, self.text_features_kmeans]: w.setToolTip('') @@ -1065,6 +1088,10 @@ def _add_connections(self): self, 'add_seg', self.check_add_seg.isChecked())) self.check_add_probas.stateChanged.connect(lambda: setattr( self, 'add_probas', self.check_add_probas.isChecked())) + self.check_add_instances.stateChanged.connect(lambda: setattr( + self, 'add_instances', self.check_add_instances.isChecked())) + self.text_inst_min_size.textChanged.connect(lambda: setattr( + self, 'inst_min_size', self.text_inst_min_size.text())) # Textboxes for PCA and Kmeans self.text_features_pca.textChanged.connect(lambda: setattr( @@ -1613,6 +1640,7 @@ def _on_select_layer(self, newtext=None): # Set flags that new outputs need to be generated (and not planes of the old ones populated) self.new_seg = True self.new_proba = True + self.new_instances = True self.new_features = True # Activate button to add annotations and segmentation layers self.add_layers_btn.setEnabled(True) @@ -1776,8 +1804,8 @@ def _on_predict(self, event=None): """Predict the segmentation of the currently viewed frame based on a classifier trained with annotations.""" - if not (self.add_seg or self.add_probas): - warnings.warn('Neither segmentation nor probabilities output selected to be added. Nothing to do.') + if not (self.add_seg or self.add_probas or self.add_instances): + warnings.warn('Neither segmentation, probabilities nor instances output selected to be added. Nothing to do.') return with warnings.catch_warnings(): @@ -1797,10 +1825,19 @@ def _on_predict(self, event=None): # Get the data image_plane = self._get_current_plane_norm() in_channels = self._parse_in_channels(self.input_channels) - - # Predict image (use backend function which returns probabilities and segmentation); skip norm as it is done above - probas, segmentation = self.cp_model._predict(image_plane, add_seg=True, in_channels=in_channels, skip_norm=True, - use_dask=self.use_dask, fe_use_device=self.fe_device) + min_size = self._parse_inst_min_size() + + # Predict image outputs; skip norm as it is done above + if self.add_seg or self.add_probas: + probas, segmentation = self.cp_model._predict(image_plane, add_seg=True, in_channels=in_channels, skip_norm=True, + use_dask=self.use_dask, fe_use_device=self.fe_device) + if self.add_instances: # If we already have segmentation, we can create instances from it directly instead of predicting them again + from napari_convpaint.utils import create_instances_from_semantic + instances = create_instances_from_semantic(segmentation, min_size=min_size) + elif self.add_instances: # If we only want instances, we can predict them directly + instances = self.cp_model.get_instances(image_plane, in_channels=in_channels, skip_norm=True, + use_dask=self.use_dask, fe_use_device=self.fe_device, + min_size=min_size) with warnings.catch_warnings(): warnings.simplefilter(action="ignore", category=FutureWarning) @@ -1840,6 +1877,21 @@ def _on_predict(self, event=None): # Case `data_dims is None` and other invalid cases are already caught above, so we don't need an else statement here self.viewer.layers[self.proba_prefix].refresh() + # Add instances if enabled + if self.add_instances: + # Check if we need to create a new instances layer + self._check_create_instances_layer() + # Set the flag to False, so we don't create a new layer every time + self.new_instances = False + + # Update instances layer + if data_dims in ['2D', '2D_RGB', '3D_multi']: + self.viewer.layers[self.instances_prefix].data = instances + elif data_dims in ['3D_single', '4D', '3D_RGB']: + self.viewer.layers[self.instances_prefix].data[step] = instances + # Case `data_dims is None` and other invalid cases are already caught above, so we don't need an else statement here + self.viewer.layers[self.instances_prefix].refresh() + def _on_get_feature_image(self, event=None): """Get the feature image for the currently viewed frame based on the current feature extractor and show it in a new layer.""" @@ -1926,9 +1978,19 @@ def _on_predict_all(self): # Predict the current step; skip normalization as it is done above in_channels = self._parse_in_channels(self.input_channels) + min_size = self._parse_inst_min_size() + # Use the backend function which returns probabilities and segmentation - probas, seg = self.cp_model._predict(image, add_seg=True, in_channels=in_channels, skip_norm=True, - use_dask=self.use_dask, fe_use_device=self.fe_device) + if self.add_seg or self.add_probas: + probas, segmentation = self.cp_model._predict(image, add_seg=True, in_channels=in_channels, skip_norm=True, + use_dask=self.use_dask, fe_use_device=self.fe_device) + if self.add_instances: # If we already have segmentation, we can create instances from it directly instead of predicting them again + from napari_convpaint.utils import create_instances_from_semantic + instances = create_instances_from_semantic(segmentation, min_size=min_size) + elif self.add_instances: # If we only want instances, we can predict them directly + instances = self.cp_model.get_instances(image, in_channels=in_channels, skip_norm=True, + use_dask=self.use_dask, fe_use_device=self.fe_device, + min_size=min_size) # In the first iteration, check if we need to create a new probas layer # (we need the information about the number of classes) @@ -1941,11 +2003,17 @@ def _on_predict_all(self): # Add the slices to the segmentation and probabilities layers if self.add_seg: - self.viewer.layers[self.seg_tag].data[step] = seg + self.viewer.layers[self.seg_tag].data[step] = segmentation self.viewer.layers[self.seg_tag].refresh() if self.add_probas: self.viewer.layers[self.proba_prefix].data[..., step, :, :] = probas self.viewer.layers[self.proba_prefix].refresh() + if self.add_instances: + if step == 0: + self._check_create_instances_layer() + self.new_instances = False + self.viewer.layers[self.instances_prefix].data[step] = instances + self.viewer.layers[self.instances_prefix].refresh() with warnings.catch_warnings(): warnings.simplefilter(action="ignore", category=FutureWarning) @@ -2149,6 +2217,7 @@ def _on_channel_mode_changed(self): # Set flags that new outputs need to be generated (and not planes of the old ones populated) self.new_seg = True self.new_proba = True + self.new_instances = True self.new_features = True def _on_norm_changed(self): @@ -2200,6 +2269,8 @@ def _on_reset_convpaint(self, event=None): self.text_input_channels.setText(self.input_channels) self.check_add_seg.setChecked(self.add_seg) self.check_add_probas.setChecked(self.add_probas) + self.check_add_instances.setChecked(self.add_instances) + self.text_inst_min_size.setText(self.inst_min_size) self.text_features_pca.setText(self.features_pca_components) self.text_features_kmeans.setText(self.features_kmeans_clusters) # Reset the model description @@ -2267,6 +2338,7 @@ def _reset_attributes(self): self.old_annot_tag = "None" # Tag for the annotations, saved to be able to rename them later self.old_seg_tag = "None" # Tag for the segmentation, saved to be able to rename them later self.old_proba_tag = "None" # Tag for the probabilities, saved to be able to rename them later + self.old_instances_tag = "None" # Tag for the instances, saved to be able to rename them later self.add_layers_flag = True # Flag to prevent adding layers twice on one trigger # self.update_layer_flag = True # Flag to prevent updating layers twice on one trigger # self.rgb_img = getattr(self, "rgb_img", None) or False # Tag to register if the image is RGB @@ -2279,6 +2351,7 @@ def _reset_attributes(self): self.annot_tag = 'annotations' # Prefix for the annotations layer names self.seg_tag = 'segmentation' # Prefix for the segmentation layer names self.proba_prefix = 'probabilities' # Prefix for the class probabilities layer names + self.instances_prefix = 'instances' # Prefix for the instances layer names self.features_prefix = 'features' # Prefix for the feature image layer name self.cont_training = "Image" # Update features for subsequent training ("Image" or "Off" or "Global") self.use_dask = False # Use Dask for parallel processing @@ -2287,8 +2360,11 @@ def _reset_attributes(self): self.input_channels = "" # Input channels for the model (as txt, will be parsed) self.add_seg = True # Add a layer with segmentation self.add_probas = False # Add a layer with class probabilities + self.add_instances = False # Add a layer with instances + self.inst_min_size = "100" # Minimum number of pixels for instances (0 = ignore) self.new_seg = True # Flags to indicate if new outputs are created self.new_proba = True + self.new_instances = True self.new_features = True self.features_pca_components = "0" # Number of PCA components for feature image (0 = no PCA) self.features_kmeans_clusters = "0" # Number of k-means clusters for feature image (0 = no k-means) @@ -2709,6 +2785,34 @@ def _check_create_probas_layer(self, num_classes): # Save information about the probabilities layer to be able to rename it later self._set_old_proba_tag() + def _check_create_instances_layer(self): + """Check if instances layer exists and create it if not.""" + + img = self._get_selected_img(check=True) + if img is None: + warnings.warn('No image selected. No layers added.') + return + + layer_shape = self._get_annot_shape(img) + num_spatial = len(layer_shape) + transform_kwargs = self._get_layer_transform_kwargs(img, num_spatial_dims=num_spatial, num_leading_dims=0) + + instances_exists = self.instances_prefix in self.viewer.layers + + if self.new_instances & instances_exists: + if self.keep_layers: + self._rename_instances_for_backup() + else: + self.viewer.layers.remove(self.instances_prefix) + + if (not instances_exists) or self.new_instances: + self.viewer.add_labels( + data=np.zeros((layer_shape), dtype=np.int32), + name=self.instances_prefix, + **transform_kwargs + ) + self._set_old_instances_tag() + def _check_create_features_layer(self, num_features): """Check if feature image layer exists and create it if not.""" @@ -2789,6 +2893,14 @@ def _rename_probas_for_backup(self): full_name = self._get_unique_layer_name(self.old_proba_tag) self.viewer.layers[self.proba_prefix].name = full_name + def _rename_instances_for_backup(self): + """Name the instances layer with a unique name according to its image, + so it can be kept when adding a new with the standard name.""" + # Rename the layer to avoid overwriting it + if self.instances_prefix in self.viewer.layers: + full_name = self._get_unique_layer_name(self.old_instances_tag) + self.viewer.layers[self.instances_prefix].name = full_name + def _rename_features_for_backup(self): """Name the features layer with a unique name according to its image, so it can be kept when adding a new with the standard name.""" @@ -2838,6 +2950,11 @@ def _set_old_proba_tag(self): This is used to rename old probabilities layers when creating new ones.""" self.old_proba_tag = f"{self.proba_prefix}_{self._get_old_data_tag()}" + def _set_old_instances_tag(self): + """Set the old instances tag based on the current image layer and data dimensions. + This is used to rename old instances layers when creating new ones.""" + self.old_instances_tag = f"{self.instances_prefix}_{self._get_old_data_tag()}" + def _set_old_features_tag(self): """Set the old features tag based on the current image layer and data dimensions. This is used to rename old features layers when creating new ones.""" @@ -3849,6 +3966,15 @@ def _parse_in_channels(channels_text): return channels except ValueError: return None + + def _parse_inst_min_size(self): + """Parse the minimum instance size from text.""" + if not self.inst_min_size: + self.inst_min_size = '100' + elif not self.inst_min_size.isdigit(): + warnings.warn('Instance min_num_pix must be an integer. Using 100.') + self.inst_min_size = '100' + return int(self.inst_min_size) ### MULTIFILE TAB @@ -4358,11 +4484,11 @@ def _on_segment_selected_multifile(self): warnings.warn('No trained model available for segmentation.') return - # If the user has selected probabilities as outputs, notify them that this is not (yet) supported and segmentations will be created and saved instead - if self.add_probas: - show_info('You selected class probabilities as output. This is not yet supported in the Multifile workflow. Segmentations will be created and saved instead.') + # If the user has selected probabilities and/or instances as outputs, notify them that this is not (yet) supported and segmentations will be created and saved instead + if self.add_probas or self.add_instances: + show_info('You selected class probabilities and/or instances as output. This is not yet supported in the Multifile workflow. Semantic segmentations will be created and saved instead.') elif not self.add_seg: - show_info('You have not selected segmentations or probabilities as output. Segmentations will be created and saved by default.') + show_info('You have not selected any output. Segmentations will be created and saved by default.') # Ask for output folder default_dir = str(Path(self._multifile_last_folder)) if getattr(self, '_multifile_last_folder', None) else str(Path.cwd()) diff --git a/src/napari_convpaint/utils.py b/src/napari_convpaint/utils.py index 84a6366..b468bad 100644 --- a/src/napari_convpaint/utils.py +++ b/src/napari_convpaint/utils.py @@ -66,6 +66,117 @@ def apply_kmeans_to_f_image(feature_img, n_clusters, random_state=None): return kmeans_labels +### Instance creation from semantic segmentaiton + +def create_instances_from_semantic(segmentations, min_size=300, classes=None, per_plane=False, warn=True): + """ + Create instance masks from semantic segmentation masks. + + Parameters: + ---------- + segmentations : list of np.ndarray or a single np.ndarray + List of semantic segmentation masks. Each mask should have shape (H, W) or (Z, H, W). + min_size : int + Minimum size of objects to keep. Smaller objects will be removed, and smaller holes will be filled. Use 0 to ignore. + classes : list of int, optional + List of classes to create instances for. If None, all classes in the segmentations will be used, except 1 (background). + per_plane : bool + If True, and the segmentation masks are 3D, each plane will be labeled separately. If False, the 3D mask will be labeled as a whole. + warn : bool + If True, a warning will be issued if no class 1 is found in the segmentations. + + """ + # Assure list input for segmentations + single_input = hasattr(segmentations, 'ndim') and segmentations.ndim >= 2 and not isinstance(segmentations, (list, tuple)) + if single_input: + segmentations = [segmentations] + + # Check for classes across segmentations + if classes is None: + classes = np.unique(np.concatenate([s.ravel() for s in segmentations])) + if not 1 in classes: + if warn: + warnings.warn(f"No class 1 found in segmentations. Found classes: {classes}.\n" + + f"This might be intentional. But typically class 1 is the background class, and classes > 1 are the objects of interest.\n" + + f"You can turn this warning off by setting `warn=False`.") + + from skimage.morphology import remove_small_holes, remove_small_objects + from skimage.measure import label + + min_distance = max(2, 2*int((min_size / np.pi)**0.5)) # Approximate min distance for ~circular objects of area min_size, to separate touching objects with watershed + + instance_masks = [np.zeros_like(segmentation, dtype=np.int32) for segmentation in segmentations] + global_max = 0 + + for c in classes: + if c == 1: + continue # Skip background class + + class_start = global_max + 1 + + for segmentation, instance_mask in zip(segmentations, instance_masks): + + semantic_class_mask = segmentation == c + + semantic_class_mask = remove_small_holes(semantic_class_mask, max_size=min_size) + semantic_class_mask = remove_small_objects(semantic_class_mask, max_size=min_size) + + if per_plane and semantic_class_mask.ndim == 3: # If we do NOT want true 3D interpretation, but have 3D masks, we loop over the planes and label them separately + labels_stack = np.zeros_like(semantic_class_mask, dtype=np.int32) + for z in range(semantic_class_mask.shape[0]): + semantic_mask_plane = semantic_class_mask[z] + instance_mask_plane = instance_mask[z] + + # from skimage.morphology import binary_erosion, disk + # semantic_mask_plane = binary_erosion(semantic_mask_plane, disk(1)) # Separate touching objects + if min_size == 0: + l = label(semantic_mask_plane) + else: + l = distance_watershed(semantic_mask_plane, min_distance=min_distance) # Use distance transform and watershed to separate touching objects + labels_stack[z] = l + + new_slice = labels_stack[z] + global_max + instance_mask_plane[semantic_mask_plane] = new_slice[semantic_mask_plane] + + global_max = instance_mask.max() + + else: # For true 3D and 2D, we can label the mask directly + # from skimage.morphology import binary_erosion, disk + # semantic_class_mask = binary_erosion(semantic_class_mask, disk(1)) # Separate touching objects + if min_size == 0: + l = label(semantic_class_mask) + else: + l = distance_watershed(semantic_class_mask, min_distance=min_distance) # Use distance transform and watershed to separate touching objects + + new_mask = l + global_max # Offset the labels by the current global max to ensure unique instance IDs across segmentations + instance_mask[semantic_class_mask] = new_mask[semantic_class_mask] + + global_max = instance_mask.max() + + + print(f"Turned class {c} into instances {class_start}-{global_max}.") + + if single_input: + return instance_masks[0] + return instance_masks + +def distance_watershed(mask, min_distance): + + from scipy import ndimage as ndi + from skimage.feature import peak_local_max + from skimage.segmentation import watershed + + distance = ndi.distance_transform_edt(mask) + + coords = peak_local_max(distance, min_distance=min_distance, labels=mask, exclude_border=False) + markers = np.zeros_like(mask, dtype=np.int32) + markers[tuple(coords.T)] = np.arange(1, len(coords)+1) + + labels = watershed(-distance, markers, mask=mask) + + return labels + + ### MODEL DOWNLOAD def guided_model_download(model_file: str, model_url: str, model_dir: str = None) -> str: