Min-Cut on Agglomerate Graph (#6361)

* [WIP] min-cut on agglomerate graph * function stubs * position2 * use mincut from jgrapht * tryo * also send segment positions of min-cut edges * add yarn kill-listeners cmd * fix merge-related compile error * pretty * add context menu entry and action for agglomerate min cut * add api method to backend route * perform min-cut in proofread saga [WIP] * pretty * pretty * weird sleep hack * perform min-cut [WIP] * pretty * remove update group assertion * remove TODOs * update unreleased Changelog * Code review changes Co-authored-by: Philipp Otto <[email protected]> Co-authored-by: Arthur Hilbert <[email protected]> Co-authored-by: Arthur Hilbert <[email protected]>
scalableminds · Sep 8, 2022 · 4f32e6f · 4f32e6f
1 parent 37571e2
commit 4f32e6f
Show file tree

Hide file tree

Showing 10 changed files with 243 additions and 22 deletions.
diff --git a/CHANGELOG.unreleased.md b/CHANGELOG.unreleased.md
@@ -13,6 +13,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 ### Added
 - Zarr-based remote dataset import now also works for public AWS S3 endpoints with no credentials. [#6421](https://github.com/scalableminds/webknossos/pull/6421)
 - Added a context menu option to extract the shortest path between two nodes as a new tree. Select the source node and open the context menu by right-clicking on another node in the same tree. [#6423](https://github.com/scalableminds/webknossos/pull/6423)
+- Added a context menu option to separate an agglomerate skeleton using Min-Cut. Activate the Proofreading tool, select the source node and open the context menu by right-clicking on the target node which you would like to separate through Min-Cut. [#6361](https://github.com/scalableminds/webknossos/pull/6361)
 
 ### Changed
 

diff --git a/frontend/javascripts/admin/admin_rest_api.ts b/frontend/javascripts/admin/admin_rest_api.ts
@@ -2271,3 +2271,24 @@ export function getSynapseTypes(
     ),
   );
 }
+
+type MinCutTargetEdge = {
+  position1: Vector3;
+  position2: Vector3;
+  segmentId1: number;
+  segmentId2: number;
+};
+export async function getEdgesForAgglomerateMinCut(
+  tracingStoreUrl: string,
+  tracingId: string,
+  segmentsInfo: Object,
+): Promise<Array<MinCutTargetEdge>> {
+  return doWithToken((token) =>
+    Request.sendJSONReceiveJSON(
+      `${tracingStoreUrl}/tracings/volume/${tracingId}/agglomerateGraphMinCut?token=${token}`,
+      {
+        data: segmentsInfo,
+      },
+    ),
+  );
+}
diff --git a/frontend/javascripts/oxalis/model/actions/skeletontracing_actions.tsx b/frontend/javascripts/oxalis/model/actions/skeletontracing_actions.tsx
@@ -41,6 +41,7 @@ type DeselectActiveTreeAction = ReturnType<typeof deselectActiveTreeAction>;
 type SetActiveGroupAction = ReturnType<typeof setActiveGroupAction>;
 type DeselectActiveGroupAction = ReturnType<typeof deselectActiveGroupAction>;
 export type MergeTreesAction = ReturnType<typeof mergeTreesAction>;
+export type MinCutAgglomerateAction = ReturnType<typeof minCutAgglomerateAction>;
 type SetTreeNameAction = ReturnType<typeof setTreeNameAction>;
 type SelectNextTreeAction = ReturnType<typeof selectNextTreeAction>;
 type SetTreeColorIndexAction = ReturnType<typeof setTreeColorIndexAction>;
@@ -82,6 +83,7 @@ export type SkeletonTracingAction =
   | SetActiveTreeByNameAction
   | DeselectActiveTreeAction
   | MergeTreesAction
+  | MinCutAgglomerateAction
   | SetTreeNameAction
   | SelectNextTreeAction
   | SetTreeColorAction
@@ -350,6 +352,13 @@ export const mergeTreesAction = (sourceNodeId: number, targetNodeId: number) =>
     targetNodeId,
   } as const);
 
+export const minCutAgglomerateAction = (sourceNodeId: number, targetNodeId: number) =>
+  ({
+    type: "MIN_CUT_AGGLOMERATE",
+    sourceNodeId,
+    targetNodeId,
+  } as const);
+
 export const setTreeNameAction = (
   name: string | undefined | null = null,
   treeId?: number | null | undefined,

diff --git a/frontend/javascripts/oxalis/model/reducers/skeletontracing_reducer_helpers.ts b/frontend/javascripts/oxalis/model/reducers/skeletontracing_reducer_helpers.ts
@@ -225,7 +225,7 @@ export function deleteEdge(
   targetTree: Tree,
   targetNode: Node,
   timestamp: number,
-): Maybe<[TreeMap, number]> {
+): Maybe<[TreeMap, number | null]> {
   return getSkeletonTracing(state.tracing).chain((skeletonTracing) => {
     if (sourceTree.treeId !== targetTree.treeId) {
       // The two selected nodes are in different trees
@@ -254,8 +254,11 @@ export function deleteEdge(
       timestamp,
     );
     // The treeId of the tree the active node belongs to could have changed
-    const newActiveTree = findTreeByNodeId(newTrees, sourceNode.id).get();
-    return Maybe.Just([newTrees, newActiveTree.treeId]);
+    const activeNodeId = skeletonTracing.activeNodeId;
+    const newActiveTreeId = activeNodeId
+      ? findTreeByNodeId(newTrees, activeNodeId).get().treeId
+      : null;
+    return Maybe.Just([newTrees, newActiveTreeId]);
   });
 }
 

diff --git a/frontend/javascripts/oxalis/model/sagas/proofread_saga.ts b/frontend/javascripts/oxalis/model/sagas/proofread_saga.ts
@@ -4,9 +4,11 @@ import { select, take } from "oxalis/model/sagas/effect-generators";
 import { AnnotationToolEnum, MappingStatusEnum, Vector3 } from "oxalis/constants";
 import Toast from "libs/toast";
 import {
+  deleteEdgeAction,
   DeleteEdgeAction,
   loadAgglomerateSkeletonAction,
   MergeTreesAction,
+  MinCutAgglomerateAction,
   setTreeNameAction,
 } from "oxalis/model/actions/skeletontracing_actions";
 import {
@@ -36,7 +38,7 @@ import {
   getMappingInfo,
   getResolutionInfo,
 } from "oxalis/model/accessors/dataset_accessor";
-import { makeMappingEditable } from "admin/admin_rest_api";
+import { getEdgesForAgglomerateMinCut, makeMappingEditable } from "admin/admin_rest_api";
 import { setMappingNameAction } from "oxalis/model/actions/settings_actions";
 import { getSegmentIdForPositionAsync } from "oxalis/controller/combinations/volume_handlers";
 import { loadAdHocMeshAction } from "oxalis/model/actions/segmentation_actions";
@@ -47,11 +49,15 @@ import { getConstructorForElementClass } from "oxalis/model/bucket_data_handling
 import { Tree } from "oxalis/store";
 import { APISegmentationLayer } from "types/api_flow_types";
 import { setBusyBlockingInfoAction } from "oxalis/model/actions/ui_actions";
+import _ from "lodash";
 
 export default function* proofreadMapping(): Saga<any> {
   yield* take("INITIALIZE_SKELETONTRACING");
   yield* take("WK_READY");
-  yield* takeEvery(["DELETE_EDGE", "MERGE_TREES"], splitOrMergeAgglomerate);
+  yield* takeEvery(
+    ["DELETE_EDGE", "MERGE_TREES", "MIN_CUT_AGGLOMERATE"],
+    splitOrMergeOrMinCutAgglomerate,
+  );
   yield* takeEvery(["PROOFREAD_AT_POSITION"], proofreadAtPosition);
 }
 
@@ -251,7 +257,15 @@ function* ensureHdf5MappingIsEnabled(layerName: string): Saga<boolean> {
   return true;
 }
 
-function* splitOrMergeAgglomerate(action: MergeTreesAction | DeleteEdgeAction) {
+let currentlyPerformingMinCut = false;
+
+function* splitOrMergeOrMinCutAgglomerate(
+  action: MergeTreesAction | DeleteEdgeAction | MinCutAgglomerateAction,
+) {
+  // Prevent this method from running recursively into itself during Min-Cut.
+  if (currentlyPerformingMinCut) {
+    return;
+  }
   const allowUpdate = yield* select((state) => state.tracing.restrictions.allowUpdate);
   if (!allowUpdate) return;
 
@@ -331,6 +345,18 @@ function* splitOrMergeAgglomerate(action: MergeTreesAction | DeleteEdgeAction) {
     agglomerateFileZoomstep,
   );
 
+  const volumeTracingWithEditableMapping = yield* select((state) =>
+    getActiveSegmentationTracing(state),
+  );
+  if (
+    volumeTracingWithEditableMapping == null ||
+    volumeTracingWithEditableMapping.mappingName == null
+  ) {
+    yield* put(setBusyBlockingInfoAction(false));
+    return;
+  }
+  const editableMappingId = volumeTracingWithEditableMapping.mappingName;
+
   /* Send the respective split/merge update action to the backend (by pushing to the save queue
      and saving immediately) */
 
@@ -364,6 +390,67 @@ function* splitOrMergeAgglomerate(action: MergeTreesAction | DeleteEdgeAction) {
         agglomerateFileMag,
       ),
     );
+  } else if (action.type === "MIN_CUT_AGGLOMERATE") {
+    if (sourceNodeAgglomerateId !== targetNodeAgglomerateId) {
+      Toast.error("Segments need to be in the same agglomerate to perform a min-cut.");
+      yield* put(setBusyBlockingInfoAction(false));
+      return;
+    }
+
+    currentlyPerformingMinCut = true;
+
+    const tracingStoreUrl = yield* select((state) => state.tracing.tracingStore.url);
+    const segmentsInfo = {
+      segmentPosition1: sourceNodePosition,
+      segmentPosition2: targetNodePosition,
+      mag: agglomerateFileMag,
+      agglomerateId: sourceNodeAgglomerateId,
+      editableMappingId,
+    };
+
+    const edgesToRemove = yield* call(
+      getEdgesForAgglomerateMinCut,
+      tracingStoreUrl,
+      volumeTracingId,
+      segmentsInfo,
+    );
+
+    for (const edge of edgesToRemove) {
+      let firstNodeId;
+      let secondNodeId;
+      for (const node of sourceTree.nodes.values()) {
+        if (_.isEqual(node.position, edge.position1)) {
+          firstNodeId = node.id;
+        } else if (_.isEqual(node.position, edge.position2)) {
+          secondNodeId = node.id;
+        }
+        if (firstNodeId && secondNodeId) {
+          break;
+        }
+      }
+
+      if (!firstNodeId || !secondNodeId) {
+        Toast.warning(
+          `Unable to find all nodes for positions ${!firstNodeId ? edge.position1 : null}${
+            !secondNodeId ? [", ", edge.position2] : null
+          } in ${sourceTree.name}.`,
+        );
+        yield* put(setBusyBlockingInfoAction(false));
+        currentlyPerformingMinCut = false;
+        return;
+      }
+
+      yield* put(deleteEdgeAction(firstNodeId, secondNodeId));
+      items.push(
+        splitAgglomerate(
+          sourceNodeAgglomerateId,
+          edge.position1,
+          edge.position2,
+          agglomerateFileMag,
+        ),
+      );
+    }
+    currentlyPerformingMinCut = false;
   }
 
   if (items.length === 0) {
@@ -378,17 +465,6 @@ function* splitOrMergeAgglomerate(action: MergeTreesAction | DeleteEdgeAction) {
 
   yield* call([api.data, api.data.reloadBuckets], layerName);
 
-  const volumeTracingWithEditableMapping = yield* select((state) =>
-    getActiveSegmentationTracing(state),
-  );
-  if (
-    volumeTracingWithEditableMapping == null ||
-    volumeTracingWithEditableMapping.mappingName == null
-  ) {
-    yield* put(setBusyBlockingInfoAction(false));
-    return;
-  }
-
   const newSourceNodeAgglomerateId = yield* call(
     [api.data, api.data.getDataValue],
     layerName,

diff --git a/frontend/javascripts/oxalis/view/context_menu.tsx b/frontend/javascripts/oxalis/view/context_menu.tsx
@@ -39,6 +39,7 @@ import {
 import {
   deleteEdgeAction,
   mergeTreesAction,
+  minCutAgglomerateAction,
   deleteNodeAction,
   setActiveNodeAction,
   createTreeAction,
@@ -98,6 +99,7 @@ type OwnProps = {
 type DispatchProps = {
   deleteEdge: (arg0: number, arg1: number) => void;
   mergeTrees: (arg0: number, arg1: number) => void;
+  minCutAgglomerate: (arg0: number, arg1: number) => void;
   deleteNode: (arg0: number, arg1: number) => void;
   setActiveNode: (arg0: number) => void;
   hideTree: (arg0: number) => void;
@@ -342,6 +344,7 @@ function NodeContextMenuOptions({
   hideContextMenu,
   deleteEdge,
   mergeTrees,
+  minCutAgglomerate,
   deleteNode,
   createBranchPoint,
   deleteBranchpointById,
@@ -352,6 +355,9 @@ function NodeContextMenuOptions({
   infoRows,
   allowUpdate,
 }: NodeContextMenuOptionsProps): JSX.Element {
+  const isProofreadingActive = useSelector(
+    (state: OxalisState) => state.uiInformation.activeTool === "PROOFREAD",
+  );
   const dispatch = useDispatch();
 
   if (skeletonTracing == null) {
@@ -401,6 +407,17 @@ function NodeContextMenuOptions({
             Create Edge & Merge with this Tree{" "}
             {useLegacyBindings ? shortcutBuilder(["Shift", "Alt", "leftMouse"]) : null}
           </Menu.Item>
+          {isProofreadingActive ? (
+            <Menu.Item
+              key="min-cut-node"
+              disabled={!areInSameTree || isTheSameNode}
+              onClick={() =>
+                activeNodeId != null ? minCutAgglomerate(clickedNodeId, activeNodeId) : null
+              }
+            >
+              Perform Min-Cut between these Nodes
+            </Menu.Item>
+          ) : null}
           <Menu.Item
             key="delete-edge"
             disabled={!areNodesConnected}
@@ -1062,6 +1079,10 @@ const mapDispatchToProps = (dispatch: Dispatch<any>) => ({
     dispatch(mergeTreesAction(sourceNodeId, targetNodeId));
   },
 
+  minCutAgglomerate(sourceNodeId: number, targetNodeId: number) {
+    dispatch(minCutAgglomerateAction(sourceNodeId, targetNodeId));
+  },
+
   deleteNode(nodeId: number, treeId: number) {
     dispatch(deleteNodeAction(nodeId, treeId));
   },

diff --git a/package.json b/package.json
@@ -102,6 +102,7 @@
     "build-dev": "node_modules/.bin/webpack",
     "build-watch": "node_modules/.bin/webpack -w",
     "listening": "lsof -i:7155,9000,9001,9002",
+    "kill-listeners": "kill $(lsof -t -i:7155,9000,9001,9002)",
     "test": "tools/test.sh test --timeout=30s",
     "test-verbose": "xvfb-run -s '-ac -screen 0 1280x1024x24' tools/test.sh test --timeout=60s --verbose",
     "test-only": "tools/test.sh test --timeout=30s",

diff --git a/...e/app/com/scalableminds/webknossos/tracingstore/controllers/VolumeTracingController.scala b/...e/app/com/scalableminds/webknossos/tracingstore/controllers/VolumeTracingController.scala
@@ -25,6 +25,7 @@ import com.scalableminds.webknossos.tracingstore.slacknotification.TSSlackNotifi
 import com.scalableminds.webknossos.tracingstore.tracings.editablemapping.{
   EditableMappingService,
   EditableMappingUpdateActionGroup,
+  MinCutParameters,
   RemoteFallbackLayer
 }
 import com.scalableminds.webknossos.tracingstore.tracings.volume.{
@@ -530,6 +531,20 @@ class VolumeTracingController @Inject()(
       }
     }
 
+  def agglomerateGraphMinCut(token: Option[String], tracingId: String): Action[MinCutParameters] =
+    Action.async(validateJson[MinCutParameters]) { implicit request =>
+      log() {
+        accessTokenService.validateAccess(UserAccessRequest.readTracing(tracingId), urlOrHeaderToken(token, request)) {
+          for {
+            tracing <- tracingService.find(tracingId)
+            _ <- bool2Fox(tracing.getMappingIsEditable) ?~> "Mapping is not editable"
+            remoteFallbackLayer <- RemoteFallbackLayer.fromVolumeTracing(tracing)
+            edges <- editableMappingService.agglomerateGraphMinCut(request.body, remoteFallbackLayer, token)
+          } yield Ok(Json.toJson(edges))
+        }
+      }
+    }
+
   def updateEditableMapping(token: Option[String], tracingId: String): Action[List[EditableMappingUpdateActionGroup]] =
     Action.async(validateJson[List[EditableMappingUpdateActionGroup]]) { implicit request =>
       accessTokenService.validateAccess(UserAccessRequest.writeTracing(tracingId), urlOrHeaderToken(token, request)) {
@@ -541,7 +556,6 @@ class VolumeTracingController @Inject()(
           _ <- bool2Fox(request.body.length == 1) ?~> "Editable mapping update group must contain exactly one update group"
           updateGroup <- request.body.headOption.toFox
           _ <- bool2Fox(updateGroup.version == currentVersion + 1) ?~> "version mismatch"
-          _ <- bool2Fox(updateGroup.actions.length == 1) ?~> "Editable mapping update group must contain exactly one update action"
           _ <- editableMappingService.update(mappingName, updateGroup, updateGroup.version)
         } yield Ok
       }