Repository: Project-MONAI/MONAILabel (language: python)

Issue:
[WIP] Implement multilabel as binary scribbles for a given segment and add support for dynamic scribbles label ids from client/server
As reported in this discussion #713 and workaround attempted in #716 the labels for scribbles are fixed which causes issues when these change in client side. In addition, as reported in the discussion that "update" from scribbles functions relies on selected label in SmartEdit tab. This need not be selected or provided explicitly for scribbles functionality to work.

To address the above, I have the following changes in this PR
- [x] Send `label_info` for any update scribbles call in Slicer plugin
- [x] Recieve and update from `labels` from scribbles server app in Slicer plugin
- [x] Read and use scribbles label indices from `label_info` in Server app
- [x] Define default labels in scribbles Server app to avoid any confusion/ambiguity when running apps with Client
- [x] Add combobox for Scribbles labels selection
- [x] Add mechanism to reset scribbles for change in label selection with a use warning
- [x] Add relevant updates from #716 (thanks adamaji!)
- [x] Test with multi-label selected
- [x] Code format
- [x] Unittests locally
- [x] Tested locally with scribbles examples


Discussion / hints:
Hi, thanks for the help here. Trying this out with the same sample app setup from #713 I'm not seeing scribbles work still. It seems like `ExportVisibleSegmentsToLabelmapNode` is still being used here, so the labels are still offset, both for scribbles and for "Submit Label". Also, is scribbles supposed to work when given a multilabel input? From https://github.com/Project-MONAI/MONAILabel/issues/607 I thought that scribbles is currently intended to be a single-class algorithm.

Would it be possible to work off of #716 instead, and to address the issue of 

>may only work for this specific case, given the user has selected required label in SmartEdit tab.

we could implement the first part of @SachidanandAlle 's suggestion here https://github.com/Project-MONAI/MONAILabel/issues/607#issuecomment-1030450954 -- which I understood to be adding a similar label combobox in the scribbles section to select the desired label to perform graphcut on?


Can you try this with a fresh setup? Please remember to also update your slicer plugin to the one in this PR. 
It should address any labels offsets, so no need to change export functions. 

Multilabel is a separate problem from this, which isn't used in the app you are testing. Let's solve one problem at a time.
Another reason I'm saying to use `ExportSegmentsToLabelmapNode` instead of `ExportVisibleSegmentsToLabelmapNode` is because the latter function makes the labels sent to the server dependent on the visibility of the segments in the MRML scene. So if you make any of the segments hidden with the "eye" button or through code, they don't get sent over. I don't think that is intended behavior, and if it is, it's unclear from the GUI that it is. `ExportSegmentsToLabelmapNode` will allow you to specify which segments specifically get sent, so you don't have to do any segment deletion either.

Additionally, I think multilabel is something that should be addressed to some degree here. I think future proofing by sending all the necessary labels is fine, but is scribbles just not supposed to work currently if you have more than 1 label class? It seems like to make things clear to the user here, we need to connect the `selected_label_name` from `self.updateSegmentationMask(result_file, [selected_label_name])` directly to the scribbles section of the plugin. This could be an additional combobox in the scribbles section.
`ExportSegmentsToLabelmapNode` is fine 👍 
possibly @adamaji and @masadcv can sync up together and put the best pieces together to make the scribbles better...
thank you both :)
> Additionally, I think multilabel is something that should be addressed to some degree here. I think future proofing by sending all the necessary labels is fine, but is scribbles just not supposed to work currently if you have more than 1 label class? It seems like to make things clear to the user here, we need to connect the `selected_label_name` from `self.updateSegmentationMask(result_file, [selected_label_name])` directly to the scribbles section of the plugin. This could be an additional combobox in the scribbles section.

Many thanks for your input and for discussion. Here we have multiple things going on, and hence I mentioned multilabel to be a separate problem. The current scribbles infer task only looks at binary problems (i.e. using graphcuts). I believe to properly address this we may want to use a different transform (e.g. CRF) that handles multilabels together. This could be an additional PR itself, so best to not touch that here. 

I agree with the suggestion that we can use `selected_label_names` in a combobox specifically for scribbles to divide this problem into multiple binary segmentation problems - instead of multilabel together. Infact I already have some update on this, will push code in a moment (it is still WIP though)

I am off now, but would be good to catchup with you next week. Please send me a message on slack (if you are there)...
Sounds good @masadcv , I'll ping you on slack next week
@SachidanandAlle This is ready for your review. 
I have added `ExportSegmentsToLabelmapNode` and syncing with @adamaji to see if there is anything left from his initial PR. 

Gave this a run with the sample setup described in https://github.com/Project-MONAI/MONAILabel/discussions/713#discussion-3978347 and saw that it works for that case. I haven't used any of the OHIF stuff so can't confirm there. Dropped some comments -- but sending the addition label->index mapping info makes sense to me, especially considering future proofing for further multi label enhancements
The OHIF scribbles seems to be broken now with this PR.  I am trying to debug it but as I have less experience with javascript, it may take a while. If someone can help me with it that would be great, otherwise I will get back to fixing it within this PR soon. 
@SachidanandAlle whats the best approach here, now that we have missing scribbles functionality? 
Shall we mark this as **not to be merged PR**?
@masadcv Isn’t the removal of SimpleCRF functionality (“These are used in scribbles as post-processing steps to refine segmentations”), just an element of Scribbles, where Scribbles can continue to exist and be improved upon independently of those post-processing steps?

Can we focus effort to solve the offset label issue as proposed in #716?
@jamesobutler Yes, you are right these are separate things. However, as you can see there are merge conflicts now due to removal of server functionality. Therefore, there is no current method that can be used to test this within monai label, hence I want to avoid merging this for now (until we have at least one method to test against). 
> @SachidanandAlle whats the best approach here, now that we have missing scribbles functionality? Shall we mark this as **not to be merged PR**?

yes.. marked it as WIP for now..
Should #716 be reopened and that incremental work continue separately?
makes sense.. because some part of the problem still needs to be solved irrespective of method (simplecrf vs others) used for scribbles..

however @masadcv is working on it.. and the problem needs to be solved at both server and client side..

Fix patch:
diff --git a/monailabel/scribbles/infer.py b/monailabel/scribbles/infer.py
new file mode 100644
index 000000000..3c1d83f77
--- /dev/null
+++ b/monailabel/scribbles/infer.py
@@ -0,0 +1,130 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from monai.transforms import Compose, EnsureChannelFirstd, LoadImaged, Orientationd, ScaleIntensityRanged, Spacingd
+
+from monailabel.interfaces.tasks.infer import InferTask, InferType
+from monailabel.scribbles.transforms import (
+    AddBackgroundScribblesFromROId,
+    ApplyGraphCutOptimisationd,
+    MakeISegUnaryd,
+    MakeLikelihoodFromScribblesHistogramd,
+)
+from monailabel.transform.post import BoundingBoxd, Restored
+
+
+class HistogramBasedGraphCut(InferTask):
+    """
+    Defines histogram-based GraphCut task for Generic segmentation from the following paper:
+
+    Wang, Guotai, et al. "Interactive medical image segmentation using deep learning with image-specific fine tuning."
+    IEEE transactions on medical imaging 37.7 (2018): 1562-1573. (preprint: https://arxiv.org/pdf/1710.04043.pdf)
+
+    This task takes as input 1) original image volume and 2) scribbles from user
+    indicating foreground and background regions. A likelihood volume is generated using histogram method.
+    User-scribbles are incorporated using Equation 7 on page 4 of the paper.
+
+    SimpleCRF's GraphCut layer is used to optimise Equation 5 from the paper, where unaries come from Equation 7
+    and pairwise is the original input volume.
+    """
+
+    def __init__(
+        self,
+        dimension=3,
+        description="A post processing step with histogram-based GraphCut for Generic segmentation",
+        intensity_range=(-300, 200, 0.0, 1.0, True),
+        pix_dim=(2.5, 2.5, 5.0),
+        lamda=1.0,
+        sigma=0.1,
+        labels=None,
+        config=None,
+    ):
+        if config:
+            config.update({"lamda": lamda, "sigma": sigma})
+        else:
+            config = {"lamda": lamda, "sigma": sigma}
+        super().__init__(
+            path=None,
+            network=None,
+            labels=labels,
+            type=InferType.SCRIBBLES,
+            dimension=dimension,
+            description=description,
+            config=config,
+        )
+        self.intensity_range = intensity_range
+        self.pix_dim = pix_dim
+        self.lamda = lamda
+        self.sigma = sigma
+
+        # set default scribbles labels
+        self.scribbles_bg_label = 2 if not self.labels else len(self.labels) + 1
+        self.scribbles_fg_label = 3 if not self.labels else len(self.labels) + 2
+
+    def pre_transforms(self, data):
+        return [
+            LoadImaged(keys=["image", "label"]),
+            EnsureChannelFirstd(keys=["image", "label"]),
+            AddBackgroundScribblesFromROId(
+                scribbles="label",
+                scribbles_bg_label=self.scribbles_bg_label,
+                scribbles_fg_label=self.scribbles_fg_label,
+            ),
+            # at the moment optimisers are bottleneck taking a long time,
+            # therefore scaling non-isotropic with big spacing
+            Spacingd(keys=["image", "label"], pixdim=self.pix_dim, mode=["bilinear", "nearest"]),
+            Orientationd(keys=["image", "label"], axcodes="RAS"),
+            ScaleIntensityRanged(
+                keys="image",
+                a_min=self.intensity_range[0],
+                a_max=self.intensity_range[1],
+                b_min=self.intensity_range[2],
+                b_max=self.intensity_range[3],
+                clip=self.intensity_range[4],
+            ),
+            MakeLikelihoodFromScribblesHistogramd(
+                image="image",
+                scribbles="label",
+                post_proc_label="prob",
+                scribbles_bg_label=self.scribbles_bg_label,
+                scribbles_fg_label=self.scribbles_fg_label,
+                normalise=True,
+            ),
+        ]
+
+    def inferer(self, data):
+        return Compose(
+            [
+                # unary term maker
+                MakeISegUnaryd(
+                    image="image",
+                    logits="prob",
+                    scribbles="label",
+                    unary="unary",
+                    scribbles_bg_label=self.scribbles_bg_label,
+                    scribbles_fg_label=self.scribbles_fg_label,
+                ),
+                # optimiser
+                ApplyGraphCutOptimisationd(
+                    unary="unary",
+                    pairwise="image",
+                    post_proc_label="pred",
+                    lamda=self.lamda,
+                    sigma=self.sigma,
+                ),
+            ]
+        )
+
+    def post_transforms(self, data):
+        return [
+            Restored(keys="pred", ref_image="image"),
+            BoundingBoxd(keys="pred", result="result", bbox="bbox"),
+        ]
diff --git a/monailabel/scribbles/transforms.py b/monailabel/scribbles/transforms.py
index 0f80fa8e3..13cca9b27 100644
--- a/monailabel/scribbles/transforms.py
+++ b/monailabel/scribbles/transforms.py
@@ -20,7 +20,7 @@
 
 from monailabel.transform.writer import Writer
 
-from .utils import make_iseg_unary, make_likelihood_image_histogram
+from .utils import make_iseg_unary, make_likelihood_image_histogram, maxflow
 
 logger = logging.getLogger(__name__)
 
@@ -64,6 +64,19 @@ def _copy_affine(self, d, src, dst):
 
         return d
 
+    def _set_scribbles_idx_from_labelinfo(self, d):
+        label_info = d.get("label_info", [])
+        for lb in label_info:
+            if lb.get("name", None) == "background_scribbles":
+                id = lb.get("id", self.scribbles_bg_label)
+                self.scribbles_bg_label = id
+                logging.info("Loading background scribbles labels from: {} with index: {}".format(lb.get("name"), id))
+
+            if lb.get("name", None) == "foreground_scribbles":
+                id = lb.get("id", self.scribbles_fg_label)
+                self.scribbles_fg_label = id
+                logging.info("Loading foreground scribbles labels from: {} with index: {}".format(lb.get("name"), id))
+
 
 #######################################
 #######################################
@@ -89,6 +102,9 @@ def __init__(
     def __call__(self, data):
         d = dict(data)
 
+        # load scribbles idx from labels_info (if available)
+        self._set_scribbles_idx_from_labelinfo(d)
+
         # read relevant terms from data
         scribbles = self._fetch_data(d, self.scribbles)
 
@@ -158,6 +174,9 @@ def __init__(
     def __call__(self, data):
         d = dict(data)
 
+        # load scribbles idx from labels_info (if available)
+        self._set_scribbles_idx_from_labelinfo(d)
+
         # copy affine meta data from image input
         d = self._copy_affine(d, src=self.image, dst=self.post_proc_label)
 
@@ -277,6 +296,9 @@ def __init__(
     def __call__(self, data):
         d = dict(data)
 
+        # load scribbles idx from labels_info (if available)
+        self._set_scribbles_idx_from_labelinfo(d)
+
         # copy affine meta data from image input
         self._copy_affine(d, self.image, self.unary)
 
@@ -311,6 +333,89 @@ def __call__(self, data):
 #######################
 #  Optimiser Transforms
 #######################
+class ApplyGraphCutOptimisationd(InteractiveSegmentationTransform):
+    """
+    Generic GraphCut optimisation transform.
+
+    This can be used in conjuction with any Make*Unaryd transform
+    (e.g. MakeISegUnaryd from above for implementing ISeg unary term).
+    It optimises a typical energy function for interactive segmentation methods using SimpleCRF's GraphCut method,
+    e.g. Equation 5 from https://arxiv.org/pdf/1710.04043.pdf.
+
+    Usage Example::
+
+        Compose(
+            [
+                # unary term maker
+                MakeISegUnaryd(
+                    image="image",
+                    logits="logits",
+                    scribbles="label",
+                    unary="unary",
+                    scribbles_bg_label=2,
+                    scribbles_fg_label=3,
+                ),
+                # optimiser
+                ApplyGraphCutOptimisationd(
+                    unary="unary",
+                    pairwise="image",
+                    post_proc_label="pred",
+                    lamda=10.0,
+                    sigma=15.0,
+                ),
+            ]
+        )
+    """
+
+    def __init__(
+        self,
+        unary: str,
+        pairwise: str,
+        meta_key_postfix: str = "meta_dict",
+        post_proc_label: str = "pred",
+        lamda: float = 8.0,
+        sigma: float = 0.1,
+    ) -> None:
+        super().__init__(meta_key_postfix)
+        self.unary = unary
+        self.pairwise = pairwise
+        self.post_proc_label = post_proc_label
+        self.lamda = lamda
+        self.sigma = sigma
+
+    def __call__(self, data):
+        d = dict(data)
+
+        # attempt to fetch algorithmic parameters from app if present
+        self.lamda = d.get("lamda", self.lamda)
+        self.sigma = d.get("sigma", self.sigma)
+
+        # copy affine meta data from pairwise input
+        self._copy_affine(d, self.pairwise, self.post_proc_label)
+
+        # read relevant terms from data
+        unary_term = self._fetch_data(d, self.unary)
+        pairwise_term = self._fetch_data(d, self.pairwise)
+
+        # check if input unary is compatible with GraphCut opt
+        if unary_term.shape[0] > 2:
+            raise ValueError(f"GraphCut can only be applied to binary probabilities, received {unary_term.shape[0]}")
+
+        # # attempt to unfold probability term
+        # unary_term = self._unfold_prob(unary_term, axis=0)
+
+        # prepare data for SimpleCRF's GraphCut
+        unary_term = torch.from_numpy(unary_term).unsqueeze(0)
+        pairwise_term = torch.from_numpy(pairwise_term).unsqueeze(0)
+
+        # run GraphCut
+        post_proc_label = maxflow(pairwise_term, unary_term, lamda=self.lamda, sigma=self.sigma).squeeze(0).numpy()
+
+        d[self.post_proc_label] = post_proc_label
+
+        return d
+
+
 class ApplyCRFOptimisationd(InteractiveSegmentationTransform):
     """
     Generic MONAI CRF optimisation transform.
diff --git a/monailabel/scribbles/utils.py b/monailabel/scribbles/utils.py
index 877e4c2df..cea85a2e4 100644
--- a/monailabel/scribbles/utils.py
+++ b/monailabel/scribbles/utils.py
@@ -11,6 +11,11 @@
 import logging
 
 import numpy as np
+from monai.utils import optional_import
+
+# torch import is needed to execute torchmaxflow
+optional_import("torch")
+import torchmaxflow
 
 logger = logging.getLogger(__name__)
 
@@ -19,6 +24,12 @@ def get_eps(data):
     return np.finfo(data.dtype).eps
 
 
+def maxflow(image, prob, lamda=5, sigma=0.1):
+    # lamda: weight of smoothing term
+    # sigma: std of intensity values
+    return torchmaxflow.maxflow(image, prob, lamda, sigma)
+
+
 def make_iseg_unary(
     prob,
     scribbles,
@@ -139,7 +150,7 @@ def make_likelihood_image_histogram(image, scrib, scribbles_bg_label, scribbles_
 
     # generate histograms for background/foreground
     bg_hist, fg_hist, bin_edges = make_histograms(
-        image, scrib, scribbles_bg_label, scribbles_fg_label, alpha_bg=1, alpha_fg=1, bins=32
+        image, scrib, scribbles_bg_label, scribbles_fg_label, alpha_bg=1, alpha_fg=1, bins=64
     )
 
     # lookup values for each voxel for generating background/foreground probabilities
diff --git a/plugins/slicer/MONAILabel/MONAILabel.py b/plugins/slicer/MONAILabel/MONAILabel.py
index 3cb51bcfb..ea402244c 100644
--- a/plugins/slicer/MONAILabel/MONAILabel.py
+++ b/plugins/slicer/MONAILabel/MONAILabel.py
@@ -222,6 +222,7 @@ def __init__(self, parent=None):
         self.timer = None
 
         self.scribblesMode = None
+        self.ignoreScribblesLabelChangeEvent = False
         self.deepedit_multi_label = False
 
     def setup(self):
@@ -289,6 +290,7 @@ def setup(self):
         self.ui.uploadImageButton.connect("clicked(bool)", self.onUploadImage)
         self.ui.importLabelButton.connect("clicked(bool)", self.onImportLabel)
         self.ui.labelComboBox.connect("currentIndexChanged(int)", self.onSelectLabel)
+        self.ui.scribLabelComboBox.connect("currentIndexChanged(int)", self.onSelectScribLabel)
         self.ui.dgUpdateButton.connect("clicked(bool)", self.onUpdateDeepgrow)
         self.ui.dgUpdateCheckBox.setStyleSheet("padding-left: 10px;")
 
@@ -311,12 +313,12 @@ def setup(self):
         self.ui.updateScribblesButton.clicked.connect(self.onUpdateScribbles)
         self.ui.paintScribblesButton.clicked.connect(self.onPaintScribbles)
         self.ui.eraseScribblesButton.clicked.connect(self.onEraseScribbles)
-        self.ui.scribblesLabelSelector.connect("currentIndexChanged(int)", self.onSelectScribblesLabel)
+        self.ui.scribblesSelector.connect("currentIndexChanged(int)", self.onSelectScribblesLabel)
 
         # creating editable combo box
-        self.ui.scribblesLabelSelector.addItem(self.icon("fg_green.png"), "Foreground")
-        self.ui.scribblesLabelSelector.addItem(self.icon("bg_red.png"), "Background")
-        self.ui.scribblesLabelSelector.setCurrentIndex(0)
+        self.ui.scribblesSelector.addItem(self.icon("fg_green.png"), "Foreground")
+        self.ui.scribblesSelector.addItem(self.icon("bg_red.png"), "Background")
+        self.ui.scribblesSelector.setCurrentIndex(0)
 
         # start with scribbles section disabled
         self.ui.scribblesCollapsibleButton.setEnabled(False)
@@ -383,7 +385,7 @@ def onSceneStartClose(self, caller, event):
             self.dgNegativePointListNodeObservers,
         )
         self.dgNegativePointListNode = None
-        self.onClearScribbles()
+        self.onResetScribbles()
 
     def resetPointList(self, markupsPlaceWidget, pointListNode, pointListNodeObservers):
         if markupsPlaceWidget.placeModeEnabled:
@@ -508,7 +510,9 @@ def updateGUIFromParameterNode(self, caller=None, event=None):
         else:
             self.ui.aclCollapsibleButton.hide()
 
+        self.ignoreScribblesLabelChangeEvent = True
         self.ui.labelComboBox.clear()
+        self.ui.scribLabelComboBox.clear()
         if self._segmentNode:
             segmentation = self._segmentNode.GetSegmentation()
             totalSegments = segmentation.GetNumberOfSegments()
@@ -516,18 +520,27 @@ def updateGUIFromParameterNode(self, caller=None, event=None):
             for idx, segmentId in enumerate(segmentIds):
                 segment = segmentation.GetSegment(segmentId)
                 label = segment.GetName()
-                if label in ["foreground_scribbles", "background_scribbles"]:
-                    continue
-                self.ui.labelComboBox.addItem(label)
+                if label not in ["foreground_scribbles", "background_scribbles"]:
+                    self.ui.labelComboBox.addItem(label)
+                if label not in ["background", "foreground_scribbles", "background_scribbles"]:
+                    self.ui.scribLabelComboBox.addItem(label)
         else:
             for label in self.info.get("labels", {}):
                 self.ui.labelComboBox.addItem(label)
+                if label != "background":
+                    self.ui.scribLabelComboBox.addItem(label)
 
         currentLabel = self._parameterNode.GetParameter("CurrentLabel")
         idx = self.ui.labelComboBox.findText(currentLabel) if currentLabel else 0
         idx = 0 if idx < 0 < self.ui.labelComboBox.count else idx
         self.ui.labelComboBox.setCurrentIndex(idx)
 
+        currentScribbleLabel = self._parameterNode.GetParameter("CurrentScribLabel")
+        idx = self.ui.scribLabelComboBox.findText(currentScribbleLabel) if currentScribbleLabel else 0
+        idx = 0 if idx < 0 < self.ui.scribLabelComboBox.count else idx
+        self.ui.scribLabelComboBox.setCurrentIndex(idx)
+        self.ignoreScribblesLabelChangeEvent = False
+
         self.ui.appComboBox.clear()
         self.ui.appComboBox.addItem(self.info.get("name", ""))
 
@@ -653,6 +666,11 @@ def updateParameterNodeFromGUI(self, caller=None, event=None):
             currentLabel = self.ui.labelComboBox.itemText(currentLabelIndex)
             self._parameterNode.SetParameter("CurrentLabel", currentLabel)
 
+        currentScribLabelIndex = self.ui.scribLabelComboBox.currentIndex
+        if currentScribLabelIndex >= 0:
+            currentScribLabel = self.ui.scribLabelComboBox.itemText(currentScribLabelIndex)
+            self._parameterNode.SetParameter("CurrentScribLabel", currentScribLabel)
+
         currentStrategyIndex = self.ui.strategyBox.currentIndex
         if currentStrategyIndex >= 0:
             currentStrategy = self.ui.strategyBox.itemText(currentStrategyIndex)
@@ -909,6 +927,14 @@ def currentSegment(self):
         logging.debug(f"Current SegmentID: {segmentId}; Segment: {segment}")
         return segmentId, segment
 
+    def currentScribSegment(self):
+        segmentation = self._segmentNode.GetSegmentation()
+        segmentId = segmentation.GetSegmentIdBySegmentName(self.ui.scribLabelComboBox.currentText)
+        segment = segmentation.GetSegment(segmentId)
+
+        logging.debug(f"Current SegmentID: {segmentId}; Segment: {segment}")
+        return segmentId, segment
+
     def onSelectLabel(self, caller=None, event=None):
         self.updateParameterNodeFromGUI(caller, event)
 
@@ -917,6 +943,22 @@ def onSelectLabel(self, caller=None, event=None):
         self.onEditControlPoints(self.dgNegativePointListNode, "MONAILabel.BackgroundPoints")
         self.ignorePointListNodeAddEvent = False
 
+    def onSelectScribLabel(self, caller=None, event=None):
+        if self.scribblesLayersPresent() and not self.ignoreScribblesLabelChangeEvent:
+            if not slicer.util.confirmOkCancelDisplay(
+                "This will clear current scribbles session.\n" "Are you sure to continue?"
+            ):
+                # undo changes to combobox
+                currentScribLabel = self._parameterNode.GetParameter("CurrentScribLabel")
+                logging.info(f"Cancel: reverting to original selection {currentScribLabel}")
+                self.ignoreScribblesLabelChangeEvent = True
+                self.ui.scribLabelComboBox.setCurrentIndex(self.ui.scribLabelComboBox.findText(currentScribLabel))
+                self.ignoreScribblesLabelChangeEvent = False
+                return
+            self.onClearScribbles()
+
+        self.updateParameterNodeFromGUI(caller, event)
+
     def icon(self, name="MONAILabel.png"):
         # It should not be necessary to modify this method
         iconPath = os.path.join(os.path.dirname(__file__), "Resources", "Icons", name)
@@ -1115,7 +1157,7 @@ def onNextSampleButton(self):
                 "Are you sure to continue?"
             ):
                 return
-            self.onClearScribbles()
+            self.onResetScribbles()
             slicer.mrmlScene.Clear(0)
 
         start = time.time()
@@ -1278,7 +1320,7 @@ def onSaveLabel(self):
         start = time.time()
         labelmapVolumeNode = None
         result = None
-        self.onClearScribbles()
+        self.onResetScribbles()
 
         if self.current_sample.get("session"):
             if not self.onUploadImage(init_sample=False):
@@ -1287,26 +1329,29 @@ def onSaveLabel(self):
         try:
             qt.QApplication.setOverrideCursor(qt.Qt.WaitCursor)
             segmentationNode = self._segmentNode
-            labelmapVolumeNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLLabelMapVolumeNode")
-            slicer.modules.segmentations.logic().ExportVisibleSegmentsToLabelmapNode(
-                segmentationNode, labelmapVolumeNode, self._volumeNode
-            )
-
             segmentation = segmentationNode.GetSegmentation()
             totalSegments = segmentation.GetNumberOfSegments()
             segmentIds = [segmentation.GetNthSegmentID(i) for i in range(totalSegments)]
 
+            # remove background and scribbles labels
             label_info = []
+            save_segment_ids = vtk.vtkStringArray()
             for idx, segmentId in enumerate(segmentIds):
                 segment = segmentation.GetSegment(segmentId)
-                if segment.GetName() in ["foreground_scribbles", "background_scribbles"]:
+                if segment.GetName() in ["background", "foreground_scribbles", "background_scribbles"]:
                     logging.info(f"Removing segment {segmentId}: {segment.GetName()}")
-                    segmentationNode.RemoveSegment(segmentId)
                     continue
 
+                save_segment_ids.InsertNextValue(segmentId)
                 label_info.append({"name": segment.GetName(), "idx": idx + 1})
                 # label_info.append({"color": segment.GetColor()})
 
+            # export labelmap
+            labelmapVolumeNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLLabelMapVolumeNode")
+            slicer.modules.segmentations.logic().ExportSegmentsToLabelmapNode(
+                segmentationNode, save_segment_ids, labelmapVolumeNode, self._volumeNode
+            )
+
             label_in = tempfile.NamedTemporaryFile(suffix=self.file_ext, dir=self.tmpdir).name
             self.reportProgress(5)
 
@@ -1732,9 +1777,20 @@ def onUpdateScribbles(self):
             # get scribbles + label
             segmentationNode = self._segmentNode
             labelmapVolumeNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLLabelMapVolumeNode")
-            slicer.modules.segmentations.logic().ExportVisibleSegmentsToLabelmapNode(
-                segmentationNode, labelmapVolumeNode, self._volumeNode
+            save_segment_ids = vtk.vtkStringArray()
+            segmentationNode.GetSegmentation().GetSegmentIDs(save_segment_ids)
+            slicer.modules.segmentations.logic().ExportSegmentsToLabelmapNode(
+                segmentationNode, save_segment_ids, labelmapVolumeNode, self._volumeNode
             )
+            segmentation = segmentationNode.GetSegmentation()
+            totalSegments = segmentation.GetNumberOfSegments()
+            segmentIds = [segmentation.GetNthSegmentID(i) for i in range(totalSegments)]
+
+            label_info = []
+            for idx, segmentId in enumerate(segmentIds):
+                segment = segmentation.GetSegment(segmentId)
+                label_info.append({"name": segment.GetName(), "id": idx + 1})
+
             scribbles_in = tempfile.NamedTemporaryFile(suffix=self.file_ext, dir=self.tmpdir).name
             self.reportProgress(5)
 
@@ -1755,6 +1811,10 @@ def onUpdateScribbles(self):
             # send scribbles + label to server along with selected scribbles method
             params = self.getParamsFromConfig("infer", scribblesMethod)
             params.update({"roi": selected_roi})
+            params.update({"label_info": label_info})
+            _, segment = self.currentScribSegment()
+            selected_label_name = segment.GetName()
+            params.update({"selected_label_name": selected_label_name})
 
             image_file = self.current_sample["id"]
             result_file, params = self.logic.infer(
@@ -1763,9 +1823,7 @@ def onUpdateScribbles(self):
 
             # display result from server
             self.reportProgress(90)
-            _, segment = self.currentSegment()
-            label = segment.GetName()
-            self.updateSegmentationMask(result_file, [label])
+            self.updateSegmentationMask(result_file, [selected_label_name])
         except:
             slicer.util.errorDisplay(
                 f"Failed to post process label on MONAI Label Server using {scribblesMethod}",
@@ -1841,7 +1899,23 @@ def onClearScribblesSegmentNodes(self):
                     labelMapRep, segmentation, segmentId, slicer.vtkSlicerSegmentationsModuleLogic.MODE_REPLACE
                 )
 
+        # refresh segmentation view to clear scribbles segmentations
+        # help from: https://discourse.slicer.org/t/refresh-volume-rendering/11847/6
+        segmentation.SetDisplayVisibility(False)
+        segmentation.SetDisplayVisibility(True)
+
     def onClearScribbles(self):
+        # for clearing scribbles and resetting tools to default
+        # remove "scribbles" segments from label
+        self.onClearScribblesSegmentNodes()
+
+        self.ui.paintScribblesButton.setChecked(True)
+        self.ui.eraseScribblesButton.setChecked(False)
+
+        self.ui.scribblesSelector.setCurrentIndex(0)
+
+    def onResetScribbles(self):
+
         # reset scribbles mode
         self.scribblesMode = None
 
@@ -1851,13 +1925,13 @@ def onClearScribbles(self):
         # remove "scribbles" segments from label
         self.onClearScribblesSegmentNodes()
 
-        # reset UI elements associated with scribbles
-        self.ui.scribblesCollapsibleButton.collapsed = True
-
         self.ui.paintScribblesButton.setChecked(False)
         self.ui.eraseScribblesButton.setChecked(False)
 
-        self.ui.scribblesLabelSelector.setCurrentIndex(0)
+        self.ui.scribblesSelector.setCurrentIndex(0)
+        self.ignoreScribblesLabelChangeEvent = True
+        self.ui.scribLabelComboBox.setCurrentIndex(0)
+        self.ignoreScribblesLabelChangeEvent = False
 
     def checkAndInitialiseScribbles(self):
         if not self._segmentNode:
@@ -1931,9 +2005,9 @@ def onSelectScribblesLabel(self):
         if not self._segmentNode:
             return
 
-        index = self.ui.scribblesLabelSelector.currentIndex
+        index = self.ui.scribblesSelector.currentIndex
         index = 0 if index < 0 else index
-        selected = self.ui.scribblesLabelSelector.itemText(index)
+        selected = self.ui.scribblesSelector.itemText(index)
 
         layer = "foreground_scribbles" if selected == "Foreground" else "background_scribbles"
         self.changeScribblesMode(layer=layer)
diff --git a/plugins/slicer/MONAILabel/Resources/UI/MONAILabel.ui b/plugins/slicer/MONAILabel/Resources/UI/MONAILabel.ui
index 031e3ed00..53f9ee102 100644
--- a/plugins/slicer/MONAILabel/Resources/UI/MONAILabel.ui
+++ b/plugins/slicer/MONAILabel/Resources/UI/MONAILabel.ui
@@ -513,15 +513,12 @@
       </property>
       <item row="2" column="1" colspan="2">
        <layout class="QGridLayout" name="gridLayout_7">
-        <item row="0" column="6">
+        <item row="4" column="6">
          <layout class="QGridLayout" name="gridLayout_8">
-          <item row="0" column="3">
-           <widget class="QPushButton" name="eraseScribblesButton">
+          <item row="0" column="4">
+           <widget class="QPushButton" name="updateScribblesButton">
             <property name="text">
-             <string>Erase</string>
-            </property>
-            <property name="checkable">
-             <bool>true</bool>
+             <string>Update</string>
             </property>
            </widget>
           </item>
@@ -535,43 +532,19 @@
             </property>
            </widget>
           </item>
-          <item row="0" column="4">
-           <widget class="QPushButton" name="updateScribblesButton">
+          <item row="0" column="3">
+           <widget class="QPushButton" name="eraseScribblesButton">
             <property name="text">
-             <string>Update</string>
+             <string>Erase</string>
+            </property>
+            <property name="checkable">
+             <bool>true</bool>
             </property>
            </widget>
           </item>
          </layout>
         </item>
-        <item row="0" column="2">
-         <widget class="QLabel" name="scribblesMethodLabel">
-          <property name="layoutDirection">
-           <enum>Qt::LeftToRight</enum>
-          </property>
-          <property name="text">
-           <string>Model:</string>
-          </property>
-         </widget>
-        </item>
-        <item row="0" column="3">
-         <widget class="QComboBox" name="scribblesMethodSelector">
-          <property name="sizePolicy">
-           <sizepolicy hsizetype="Expanding" vsizetype="Fixed">
-            <horstretch>0</horstretch>
-            <verstretch>0</verstretch>
-           </sizepolicy>
-          </property>
-         </widget>
-        </item>
-        <item row="1" column="2" rowspan="2">
-         <widget class="QLabel" name="label_13">
-          <property name="text">
-           <string>Label:</string>
-          </property>
-         </widget>
-        </item>
-        <item row="1" column="6">
+        <item row="2" column="6">
          <layout class="QHBoxLayout" name="horizontalLayout">
           <item>
            <widget class="QLabel" name="label_11">
@@ -623,19 +596,52 @@
           </item>
          </layout>
         </item>
-        <item row="1" column="3">
-         <layout class="QGridLayout" name="gridLayout_11">
-          <item row="0" column="0">
-           <widget class="QComboBox" name="scribblesLabelSelector">
-            <property name="sizePolicy">
-             <sizepolicy hsizetype="Expanding" vsizetype="Fixed">
-              <horstretch>0</horstretch>
-              <verstretch>0</verstretch>
-             </sizepolicy>
-            </property>
-           </widget>
-          </item>
-         </layout>
+        <item row="4" column="2">
+         <widget class="QLabel" name="label_2">
+          <property name="text">
+           <string>Scribbles:</string>
+          </property>
+         </widget>
+        </item>
+        <item row="2" column="2" rowspan="2">
+         <widget class="QLabel" name="label_13">
+          <property name="text">
+           <string>Label:</string>
+          </property>
+         </widget>
+        </item>
+        <item row="0" column="3">
+         <widget class="QComboBox" name="scribblesMethodSelector">
+          <property name="sizePolicy">
+           <sizepolicy hsizetype="Expanding" vsizetype="Fixed">
+            <horstretch>0</horstretch>
+            <verstretch>0</verstretch>
+           </sizepolicy>
+          </property>
+         </widget>
+        </item>
+        <item row="0" column="2">
+         <widget class="QLabel" name="scribblesMethodLabel">
+          <property name="layoutDirection">
+           <enum>Qt::LeftToRight</enum>
+          </property>
+          <property name="text">
+           <string>Model:</string>
+          </property>
+         </widget>
+        </item>
+        <item row="4" column="3">
+         <widget class="QComboBox" name="scribblesSelector">
+          <property name="sizePolicy">
+           <sizepolicy hsizetype="Expanding" vsizetype="Fixed">
+            <horstretch>0</horstretch>
+            <verstretch>0</verstretch>
+           </sizepolicy>
+          </property>
+         </widget>
+        </item>
+        <item row="2" column="3">
+         <widget class="QComboBox" name="scribLabelComboBox"/>
         </item>
        </layout>
       </item>
diff --git a/requirements.txt b/requirements.txt
index fd40a7b99..70ac88607 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -26,5 +26,6 @@ openslide-python==1.1.2
 opencv-python-headless==4.5.5.64
 Shapely==1.8.1.post1
 girder_client==3.1.8
+torchmaxflow==0.0.6rc1
 
 #sudo apt-get install openslide-tools -y
diff --git a/sample-apps/radiology/main.py b/sample-apps/radiology/main.py
index 0d2bae73a..8f67bbadf 100644
--- a/sample-apps/radiology/main.py
+++ b/sample-apps/radiology/main.py
@@ -24,6 +24,7 @@
 from monailabel.interfaces.tasks.scoring import ScoringMethod
 from monailabel.interfaces.tasks.strategy import Strategy
 from monailabel.interfaces.tasks.train import TrainTask
+from monailabel.scribbles.infer import HistogramBasedGraphCut
 from monailabel.tasks.activelearning.random import Random
 from monailabel.tasks.infer.deepgrow_pipeline import InferDeepgrowPipeline
 from monailabel.utils.others.class_utils import get_class_names
@@ -111,6 +112,21 @@ def init_infers(self) -> Dict[str, InferTask]:
                 logger.info(f"+++ Adding Inferer:: {k} => {v}")
                 infers[k] = v
 
+        #################################################
+        # Scribbles
+        #################################################
+        infers.update(
+            {
+                "Histogram+GraphCut": HistogramBasedGraphCut(
+                    intensity_range=(-300, 200, 0.0, 1.0, True),
+                    pix_dim=(2.5, 2.5, 5.0),
+                    lamda=1.0,
+                    sigma=0.1,
+                    labels=task_config.labels,
+                ),
+            }
+        )
+
         #################################################
         # Pipeline based on existing infers
         #################################################
diff --git a/setup.cfg b/setup.cfg
index fb672b02d..6f9ab590b 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -52,6 +52,7 @@ install_requires =
     opencv-python-headless==4.5.5.64
     Shapely==1.8.1.post1
     girder_client==3.1.8
+    torchmaxflow==0.0.6rc1
 
 [flake8]
 select = B,C,E,F,N,P,T4,W,B9


Test patch:
diff --git a/tests/unit/scribbles/test_transforms_infer.py b/tests/unit/scribbles/test_transforms_infer.py
index e4271679b..61b7cd0ff 100644
--- a/tests/unit/scribbles/test_transforms_infer.py
+++ b/tests/unit/scribbles/test_transforms_infer.py
@@ -13,17 +13,21 @@
 import unittest
 
 import numpy as np
+from monai.transforms import LoadImage
 from monai.utils import set_determinism
 from parameterized import parameterized
 
+from monailabel.scribbles.infer import HistogramBasedGraphCut
 from monailabel.scribbles.transforms import (
     AddBackgroundScribblesFromROId,
+    ApplyGraphCutOptimisationd,
     InteractiveSegmentationTransform,
     MakeISegUnaryd,
     MakeLikelihoodFromScribblesHistogramd,
     SoftenProbSoftmax,
     WriteLogits,
 )
+from monailabel.transform.writer import Writer
 
 set_determinism(seed=123)
 
@@ -290,6 +294,21 @@ def test_add_bg_roi_transform(self, input_param, test_input, expected_shape):
         self.assertTupleEqual(expected_shape, result["scribbles"].shape)
         self.assertTupleEqual(test_input["scribbles"].shape, result["scribbles"].shape)
 
+    @parameterized.expand(TEST_CASE_OPTIM_TX)
+    def test_optimisation_transforms(self, input_param, test_input, output, expected_shape):
+        input_param.update({"post_proc_label": "pred"})
+        for current_tx in [ApplyGraphCutOptimisationd]:
+            result = current_tx(**input_param)(test_input)
+
+            # removed assert_equal as this is non-deterministic func
+            # np.testing.assert_equal(output["target"], result["pred"])
+
+            self.assertTupleEqual(expected_shape, result["pred"].shape)
+
+        with self.assertRaises(ValueError):
+            test_input["prob"] = np.random.rand(3, 128, 128, 128)
+            result = ApplyGraphCutOptimisationd(**input_param)(test_input)
+
     @parameterized.expand(TEST_CASE_MAKE_ISEG_UNARY_TX)
     def test_make_iseg_unary_transform(self, input_param, test_input, output, expected_shape):
         input_param.update({"unary": "pred"})
@@ -361,5 +380,29 @@ def test_soften_prob_softmax_transforms(self):
         self.assertEqual(output_data["prob"].shape, data["logits"].shape)
 
 
+class TestScribblesInferers(unittest.TestCase):
+    @parameterized.expand(TEST_CASE_HISTOGRAM_GRAPHCUT)
+    def test_histogram_graphcut_inferer(self, test_input, expected_shape):
+        test_input.update({"image_path": "fakepath.nii"})
+
+        # save data to file and update test dictionary
+        image_file, data = Writer(label="image", nibabel=True)(test_input)
+        scribbles_file, data = Writer(label="label", nibabel=True)(test_input)
+
+        # add paths to file, remove any associated meta_dict
+        test_input["image"] = image_file
+        test_input["label"] = scribbles_file
+
+        test_input.pop("image_meta_dict", None)
+        test_input.pop("label_meta_dict", None)
+
+        # run scribbles inferer and load results
+        result_file, _ = HistogramBasedGraphCut()(test_input)
+        result = LoadImage()(result_file)[0]
+
+        # can only check output shape due to non-deterministic results
+        self.assertTupleEqual(expected_shape, result.shape)
+
+
 if __name__ == "__main__":
     unittest.main()