Merge branch 'master' into fixHalo

qin-yu · web-flow · commit c5ec9ee55f98 · 2024-03-11T16:25:59.000+01:00
diff --git a/README.md b/README.md
@@ -301,77 +301,8 @@ the boundary signal, we could leverage the fact that one cell contains only one
 segmentation strategy and obtain improved segmentation. This workflow is now available in all PlantSeg interfaces.
 
 ## Troubleshooting
-* If you stumble in the following error message:
-```
-AssertionError:
-The NVIDIA driver on your system is too old (found version xxxx).
-Please update your GPU driver by downloading and installing a new
-version from the URL: http://www.nvidia.com/Download/index.aspx
-Alternatively, go to: http://pytorch.org to install
-a PyTorch version that has been compiled with your version
-of the CUDA driver.
-```
-or:
-```
-    raise AssertionError("Torch not compiled with CUDA enabled")
-AssertionError: Torch not compiled with CUDA enabled
-```
-It means that your cuda installation does not match the default in plantseg.
-You can check your current cuda version by typing in the terminal
-```
-cat /usr/local/cuda/version.txt
-```
-Then you can re-install the pytorch version compatible with your cuda by activating your `plant-seg` environment:
-```
-conda activate plant-seg
-```
-and
-```
-conda install -c pytorch torchvision cudatoolkit=<YOU_CUDA_VERSION> pytorch
-```
-e.g. for cuda 9.2
-```
-conda install -c pytorch torchvision cudatoolkit=9.2 pytorch
-```
-
-Alternatively one can create the `plant-seg` environment from scratch and ensuring the correct version of cuda/pytorch, by:
-```
-conda create -n plant-seg -c lcerrone -c abailoni -c cpape -c conda-forge cudatoolkit=<YOU_CUDA_VERSION> plantseg
-```
-
-* If you use plantseg from the GUI and you receive an error similar to:
-```
-RuntimeError: key : 'crop_volume' is missing, plant-seg requires 'crop_volume' to run
-```
-(or a similar message for any of the other keys)
-It might be that the last session configuration file got corrupted or is outdated.
-You should be able to solve it by removing the corrupted file `config_gui_last.yaml`.
-
-If you have a standard installation of plantseg, you can remove it by executing on the terminal:
-```
-$ rm ~/.plantseg_models/configs/config_gui_last.yaml
-```
-
-* If you use plantseg from the command line, and you receive an error similar to:
-```
-RuntimeError: key : 'crop_volume' is missing, plant-seg requires 'crop_volume' to run
-```
-
-Please make sure that your configuration has the correct formatting and contains all required keys.
-An updated example can be found inside the directory `examples`, in this repository.
-
-* If when trying to execute the Lifted Multicut pipeline you receive an error like:
-```
-'cannot import name 'lifted_problem_from_probabilities' from 'elf.segmentation.features''
-```
-The solution is to re-install [elf](https://github.com/constantinpape/elf) via
-```
-conda install -c conda-forge python-elf
-```
 
-* PlantSeg is under active development, so it may happen that the models/configuration files saved in `~/.plantseg_modes`
-are outdated. In case of errors related to loading the configuration file, please close the PlantSeg app,
-remove `~/.plantseg_models` directory and try again.
+See [troubleshooting](https://hci-unihd.github.io/plant-seg/chapters/getting_started/troubleshooting) for a list of common issues and their solutions.
 
 ## Tests
 In order to run tests make sure that `pytest` is installed in your conda environment. You can run your tests
diff --git a/docs/chapters/getting_started/troubleshooting.md b/docs/chapters/getting_started/troubleshooting.md
@@ -1,6 +1,30 @@
-# Troubleshooting
+# Troubleshooting  <!-- omit in toc -->
 
-* If you stumble in the following error message:
+- [`Could not load library libcudnn_ops_infer.so.8.`](#could-not-load-library-libcudnn_ops_inferso8)
+- [`NVIDIA driver on your system is too old` or `Torch not compiled with CUDA enabled`](#nvidia-driver-on-your-system-is-too-old-or-torch-not-compiled-with-cuda-enabled)
+- [`RuntimeError: key : KEY_NAME is missing, plant-seg requires KEY_NAME to run`](#runtimeerror-key--key_name-is-missing-plant-seg-requires-key_name-to-run)
+- [`cannot import name 'lifted_problem_from_probabilities'`](#cannot-import-name-lifted_problem_from_probabilities)
+- [Other issues](#other-issues)
+
+----
+
+#### `Could not load library libcudnn_ops_infer.so.8.`
+
+If you stumble in the following error message:
+```
+Could not load library libcudnn_ops_infer.so.8. Error: libcudnn_ops_infer.so.8: cannot open shared object file: No such file or directory
+```
+
+Just install `cudnn` by running:
+```
+$ mamba install -c conda-forge cudnn
+```
+
+----
+
+#### `NVIDIA driver on your system is too old` or `Torch not compiled with CUDA enabled`
+
+If you stumble in the following error message:
 ```
 AssertionError:
 The NVIDIA driver on your system is too old (found version xxxx).
@@ -35,10 +59,14 @@ conda install -c pytorch torchvision cudatoolkit=9.2 pytorch
 
 Alternatively one can create the `plant-seg` environment from scratch and ensuring the correct version of cuda/pytorch, by:
 ```
-conda create -n plant-seg -c lcerrone -c abailoni -c cpape -c awolny -c conda-forge cudatoolkit=<YOU_CUDA_VERSION> plantseg
+conda create -n plant-seg -c lcerrone -c abailoni -c cpape -c conda-forge cudatoolkit=<YOU_CUDA_VERSION> plantseg
 ```
 
-* If you use plantseg from the GUI and you receive an error similar to:
+----
+
+#### `RuntimeError: key : KEY_NAME is missing, plant-seg requires KEY_NAME to run`
+
+If you use plantseg from the GUI and you receive an error similar to:
 ```
 RuntimeError: key : 'crop_volume' is missing, plant-seg requires 'crop_volume' to run
 ```
@@ -51,15 +79,19 @@ If you have a standard installation of plantseg, you can remove it by executing
 $ rm ~/.plantseg_models/configs/config_gui_last.yaml
 ```
 
-* If you use plantseg from the comand line and you receive an error similar to:
+If you use plantseg from the command line, and you receive an error similar to:
 ```
 RuntimeError: key : 'crop_volume' is missing, plant-seg requires 'crop_volume' to run
 ```
 
-Please make sure that your configuratiuon has the correct formatting and contains all required keys.
+Please make sure that your configuration has the correct formatting and contains all required keys.
 An updated example can be found inside the directory `examples`, in this repository.
 
-* If when trying to execute the Lifted Multicut pipeline you receive an error like:
+----
+
+#### `cannot import name 'lifted_problem_from_probabilities'`
+
+If when trying to execute the Lifted Multicut pipeline you receive an error like:
 ```
 'cannot import name 'lifted_problem_from_probabilities' from 'elf.segmentation.features''
 ```
@@ -68,6 +100,10 @@ The solution is to re-install [elf](https://github.com/constantinpape/elf) via
 conda install -c conda-forge python-elf
 ```
 
-* PlantSeg is under active development so it may happen that the models/configuration files saved in `~/.plantseg_modes`
+----
+
+#### Other issues
+
+* PlantSeg is under active development, so it may happen that the models/configuration files saved in `~/.plantseg_modes`
 are outdated. In case of errors related to loading the configuration file, please close the PlantSeg app,
-remove `~/.plantseg_models` directory and try again.
+remove `~/.plantseg_models` directory and try again.
diff --git a/plantseg/dataprocessing/functional/advanced_dataprocessing.py b/plantseg/dataprocessing/functional/advanced_dataprocessing.py
@@ -5,8 +5,8 @@
 import numpy as np
 import tqdm
 from skimage.filters import gaussian
-from skimage.segmentation import watershed
-
+from skimage.segmentation import watershed, relabel_sequential
+from skimage.measure import regionprops
 
 def get_bbox(mask: np.array, pixel_toll: int = 0) -> tuple[tuple, int, int, int]:
     """
@@ -235,3 +235,52 @@ def fix_over_under_segmentation_from_nuclei(cell_seg: np.array,
                                        cell_assignments,
                                        cell_idx=None)
     return _cell_seg
+
+
+def remove_false_positives_by_foreground_probability(segmentation: np.array,
+                                                     foreground: np.array,
+                                                     threshold: float) -> np.array:
+    """
+    Remove false positive regions in a segmentation based on a foreground probability map in a smart way.
+    If the mean(an instance * its own probability region) < threshold, it is removed.
+
+    Args:
+        segmentation (np.ndarray): The segmentation array, where each unique non-zero value indicates a distinct region.
+        foreground (np.ndarray): The foreground probability map, same shape as `segmentation`.
+        threshold (float): Probability threshold below which regions are considered false positives.
+
+    Returns:
+        np.ndarray: The modified segmentation array with false positives removed.
+    """
+    # TODO: make a channel for removed regions for easier inspection
+    # TODO: use `relabel_sequential` to recover the original labels
+
+    if not segmentation.shape == foreground.shape:
+        raise ValueError("Shape of segmentation and probability map must match.")
+    if foreground.max() > 1:
+        raise ValueError("Foreground must be a probability map probability map.")
+
+    instances, _, _ = relabel_sequential(segmentation)
+
+    regions = regionprops(instances)
+    to_keep = np.ones(len(regions) + 1)
+    pixel_count = np.zeros(len(regions) + 1)
+    pixel_value = np.zeros(len(regions) + 1)
+
+    for region in tqdm.tqdm(regions):
+        bbox = region.bbox
+        cube = instances[bbox[0]:bbox[3], bbox[1]:bbox[4], bbox[2]:bbox[5]] == region.label  # other instances may exist, don't use `> 0`
+        prob = foreground[bbox[0]:bbox[3], bbox[1]:bbox[4], bbox[2]:bbox[5]]
+        pixel_count[region.label] = region.area
+        pixel_value[region.label] = (cube * prob).sum()
+
+    likelihood = pixel_value / pixel_count
+    to_keep[likelihood < threshold] = 0
+    ids_to_delete = np.argwhere(to_keep == 0)
+    assert ids_to_delete.shape[1] == 1
+    ids_to_delete = ids_to_delete.flatten()
+    # print(f"    Removing instance {region.label}: pixel count: {pixel_count}, pixel value: {pixel_value}, likelihood: {likelihood}")
+
+    instances[np.isin(instances, ids_to_delete)] = 0
+    instances, _, _ = relabel_sequential(instances)
+    return instances
diff --git a/plantseg/segmentation/gasp.py b/plantseg/segmentation/gasp.py
@@ -56,9 +56,7 @@ def __init__(self,
 
         # Postprocessing size threshold
         self.post_minsize = post_minsize
-
         self.n_threads = n_threads
-
         self.dt_watershed = partial(dt_watershed,
                                     threshold=ws_threshold, sigma_seeds=ws_sigma,
                                     stacked=ws_2D, sigma_weights=ws_w_sigma,
diff --git a/plantseg/utils.py b/plantseg/utils.py
@@ -26,7 +26,7 @@ def load_config(config_path: str) -> dict:
     return config
 
 
-def get_model_zoo() -> dict:
+def get_model_zoo(get_custom: bool = True) -> dict:
     """
     returns a dictionary of all models in the model zoo.
     example:
@@ -43,22 +43,24 @@ def get_model_zoo() -> dict:
 
     zoo_config = load_config(zoo_config)
 
-    custom_zoo_config = load_config(custom_zoo)
+    if get_custom:
+        custom_zoo_config = load_config(custom_zoo)
 
-    if custom_zoo_config is None:
-        custom_zoo_config = {}
+        if custom_zoo_config is None:
+            custom_zoo_config = {}
 
-    zoo_config.update(custom_zoo_config)
+        zoo_config.update(custom_zoo_config)
     return zoo_config
 
 
 def list_models(dimensionality_filter: list[str] = None,
                 modality_filter: list[str] = None,
-                output_type_filter: list[str] = None) -> list[str]:
+                output_type_filter: list[str] = None,
+                use_custom_models: bool = True) -> list[str]:
     """
     return a list of models in the model zoo by name
     """
-    zoo_config = get_model_zoo()
+    zoo_config = get_model_zoo(use_custom_models)
     models = list(zoo_config.keys())
 
     if dimensionality_filter is not None:
diff --git a/plantseg/viewer/containers.py b/plantseg/viewer/containers.py
@@ -13,6 +13,7 @@
 from plantseg.viewer.widget.proofreading.proofreading import widget_split_and_merge_from_scribbles
 from plantseg.viewer.widget.segmentation import widget_dt_ws, widget_agglomeration
 from plantseg.viewer.widget.segmentation import widget_fix_over_under_segmentation_from_nuclei
+from plantseg.viewer.widget.segmentation import widget_fix_false_positive_from_foreground_pmap
 from plantseg.viewer.widget.segmentation import widget_lifted_multicut
 from plantseg.viewer.widget.segmentation import widget_simple_dt_ws
 
@@ -65,7 +66,8 @@ def get_gasp_workflow():
 def get_extra_seg():
     container = MainWindow(widgets=[widget_dt_ws,
                                     widget_lifted_multicut,
-                                    widget_fix_over_under_segmentation_from_nuclei],
+                                    widget_fix_over_under_segmentation_from_nuclei,
+                                    widget_fix_false_positive_from_foreground_pmap],
                            labels=False)
     container = setup_menu(container, path='https://hci-unihd.github.io/plant-seg/chapters/plantseg_interactive_napari/extra_seg.html')
     return container
diff --git a/plantseg/viewer/widget/predictions.py b/plantseg/viewer/widget/predictions.py
@@ -161,11 +161,12 @@ def _on_model_name_changed(model_name: str):
     widget_unet_predictions.model_name.tooltip = f'Select a pretrained model. Current model description: {description}'
 
 
-def _compute_multiple_predictions(image, patch_size, patch_halo, device):
+def _compute_multiple_predictions(image, patch_size, patch_halo, device, use_custom_models=True):
     out_layers = []
-    for i, model_name in enumerate(list_models()):
+    model_list = list_models(use_custom_models=use_custom_models)
+    for i, model_name in enumerate(model_list):
 
-        napari_formatted_logging(f'Running UNet Predictions: {model_name} {i}/{len(list_models())}',
+        napari_formatted_logging(f'Running UNet Predictions: {model_name} {i}/{len(model_list)}',
                                  thread='UNet Grid Predictions')
 
         out_name = create_layer_name(image.name, model_name)
@@ -193,17 +194,21 @@ def _compute_multiple_predictions(image, patch_size, patch_halo, device):
           patch_halo={'label': 'Patch halo',
                       'tooltip': 'Patch halo is extra padding for correct prediction on image boarder.'},
           device={'label': 'Device',
-                  'choices': ALL_DEVICES}
+                  'choices': ALL_DEVICES},
+          use_custom_models={'label': 'Use custom models',
+                             'tooltip': 'If True, custom models will also be used.'}
           )
 def widget_test_all_unet_predictions(image: Image,
                                      patch_size: Tuple[int, int, int] = (80, 170, 170),
                                      patch_halo: Tuple[int, int, int] = (2, 4, 4),
-                                     device: str = ALL_DEVICES[0]) -> Future[List[LayerDataTuple]]:
+                                     device: str = ALL_DEVICES[0],
+                                     use_custom_models: bool = True) -> Future[List[LayerDataTuple]]:
     func = thread_worker(partial(_compute_multiple_predictions,
                                  image=image,
                                  patch_size=patch_size,
                                  patch_halo=patch_halo,
-                                 device=device))
+                                 device=device,
+                                 use_custom_models=use_custom_models,))
 
     future = Future()
 
diff --git a/plantseg/viewer/widget/segmentation.py b/plantseg/viewer/widget/segmentation.py
@@ -7,7 +7,7 @@
 from napari.types import LayerDataTuple
 
 from napari import Viewer
-from plantseg.dataprocessing.functional.advanced_dataprocessing import fix_over_under_segmentation_from_nuclei
+from plantseg.dataprocessing.functional.advanced_dataprocessing import fix_over_under_segmentation_from_nuclei, remove_false_positives_by_foreground_probability
 from plantseg.dataprocessing.functional.dataprocessing import normalize_01
 from plantseg.segmentation.functional import gasp, multicut, dt_watershed, mutex_ws
 from plantseg.segmentation.functional import lifted_multicut_from_nuclei_segmentation, lifted_multicut_from_nuclei_pmaps
@@ -330,5 +330,36 @@ def widget_fix_over_under_segmentation_from_nuclei(cell_segmentation: Labels,
                                    input_keys=inputs_names,
                                    layer_kwarg=layer_kwargs,
                                    layer_type=layer_type,
-                                   step_name=f'Fix Over / Under segmentation',
+                                   step_name='Fix Over / Under Segmentation',
+                                   )
+
+
+@magicgui(call_button='Run Segmentation Fix from Foreground Pmap',
+          segmentation={'label': 'Segmentation'},
+          foreground={'label': 'Foreground Pmap'},
+          threshold={'label': 'Threshold',
+                     'widget_type': 'FloatSlider', 'max': 1., 'min': 0.})
+def widget_fix_false_positive_from_foreground_pmap(segmentation: Labels,
+                                                   foreground: Image,  # TODO: maybe also allow labels
+                                                   threshold=0.6) -> Future[LayerDataTuple]:
+    out_name = create_layer_name(segmentation.name, 'FGPmapFix')
+
+    inputs_names = (segmentation.name, foreground.name)
+    func_kwargs = {'segmentation': segmentation.data,
+                   'foreground': foreground.data}
+
+    layer_kwargs = layer_properties(name=out_name,
+                                    scale=segmentation.scale,
+                                    metadata=segmentation.metadata)
+    layer_type = 'labels'
+    step_kwargs = dict(threshold=threshold)
+
+    return start_threading_process(remove_false_positives_by_foreground_probability,
+                                   runtime_kwargs=func_kwargs,
+                                   statics_kwargs=step_kwargs,
+                                   out_name=out_name,
+                                   input_keys=inputs_names,
+                                   layer_kwarg=layer_kwargs,
+                                   layer_type=layer_type,
+                                   step_name='Reduce False Positives',
                                    )
