feat(sv_split): sv split in frontend

Author: akhileshh

pychunkedgraph/app/segmentation/common.py

@@ -9,16 +9,13 @@
 
 import numpy as np
 import pandas as pd
+import fastremap
 from flask import current_app, g, jsonify, make_response, request
 from pytz import UTC
 
 from pychunkedgraph import __version__
 from pychunkedgraph.app import app_utils
-from pychunkedgraph.graph import (
-    attributes,
-    cutting,
-    segmenthistory,
-)
+from pychunkedgraph.graph import attributes, cutting, segmenthistory, ChunkedGraph
 from pychunkedgraph.graph import (
     edges as cg_edges,
 )
@@ -27,6 +24,7 @@
 )
 from pychunkedgraph.graph.analysis import pathing
 from pychunkedgraph.graph.attributes import OperationLogs
+from pychunkedgraph.graph.edits_sv import split_supervoxel
 from pychunkedgraph.graph.misc import get_contact_sites
 from pychunkedgraph.graph.operation import GraphEditOperation
 from pychunkedgraph.graph.utils import basetypes
@@ -393,7 +391,7 @@ def handle_merge(table_id, allow_same_segment_merge=False):
     current_app.operation_id = ret.operation_id
     if ret.new_root_ids is None:
         raise cg_exceptions.InternalServerError(
-            "Could not merge selected " "supervoxel."
+            f"{ret.operation_id}: Could not merge selected supervoxels."
         )
 
     current_app.logger.debug(("lvl2_nodes:", ret.new_lvl2_ids))
@@ -407,24 +405,10 @@ def handle_merge(table_id, allow_same_segment_merge=False):
 ### SPLIT ----------------------------------------------------------------------
 
 
-def handle_split(table_id):
-    current_app.table_id = table_id
-    user_id = str(g.auth_user.get("id", current_app.user_id))
-
-    data = json.loads(request.data)
-    is_priority = request.args.get("priority", True, type=str2bool)
-    remesh = request.args.get("remesh", True, type=str2bool)
-    mincut = request.args.get("mincut", True, type=str2bool)
-
+def _get_sources_and_sinks(cg: ChunkedGraph, data):
     current_app.logger.debug(data)
-
-    # Call ChunkedGraph
-    cg = app_utils.get_cg(table_id, skip_cache=True)
     node_idents = []
-    node_ident_map = {
-        "sources": 0,
-        "sinks": 1,
-    }
+    node_ident_map = {"sources": 0, "sinks": 1}
     coords = []
     node_ids = []
 
@@ -437,18 +421,74 @@ def handle_split(table_id):
     node_ids = np.array(node_ids, dtype=np.uint64)
     coords = np.array(coords)
     node_idents = np.array(node_idents)
+
+    start = time.time()
     sv_ids = app_utils.handle_supervoxel_id_lookup(cg, coords, node_ids)
+    current_app.logger.info(f"SV lookup took {time.time() - start}s.")
     current_app.logger.debug(
         {"node_id": node_ids, "sv_id": sv_ids, "node_ident": node_idents}
     )
 
+    source_ids = sv_ids[node_idents == 0]
+    sink_ids = sv_ids[node_idents == 1]
+    source_coords = coords[node_idents == 0]
+    sink_coords = coords[node_idents == 1]
+    return (source_ids, sink_ids, source_coords, sink_coords)
+
+
+def handle_split(table_id):
+    current_app.table_id = table_id
+    user_id = str(g.auth_user.get("id", current_app.user_id))
+
+    data = json.loads(request.data)
+    is_priority = request.args.get("priority", True, type=str2bool)
+    remesh = request.args.get("remesh", True, type=str2bool)
+    mincut = request.args.get("mincut", True, type=str2bool)
+
+    cg = app_utils.get_cg(table_id, skip_cache=True)
+    sources, sinks, source_coords, sink_coords = _get_sources_and_sinks(cg, data)
     try:
         ret = cg.remove_edges(
             user_id=user_id,
-            source_ids=sv_ids[node_idents == 0],
-            sink_ids=sv_ids[node_idents == 1],
-            source_coords=coords[node_idents == 0],
-            sink_coords=coords[node_idents == 1],
+            source_ids=sources,
+            sink_ids=sinks,
+            source_coords=source_coords,
+            sink_coords=sink_coords,
+            mincut=mincut,
+        )
+    except cg_exceptions.SupervoxelSplitRequiredError as e:
+        current_app.logger.info(e)
+        sources_remapped = fastremap.remap(
+            sources,
+            e.sv_remapping,
+            preserve_missing_labels=True,
+            in_place=False,
+        )
+        sinks_remapped = fastremap.remap(
+            sinks,
+            e.sv_remapping,
+            preserve_missing_labels=True,
+            in_place=False,
+        )
+        overlap_mask = np.isin(sources_remapped, sinks_remapped)
+        for sv_to_split in np.unique(sources_remapped[overlap_mask]):
+            _mask0 = sources_remapped[sources_remapped == sv_to_split]
+            _mask1 = sinks_remapped[sinks_remapped == sv_to_split]
+            split_supervoxel(
+                cg,
+                sv_to_split,
+                source_coords[_mask0],
+                sink_coords[_mask1],
+                e.operation_id,
+            )
+
+        sources, sinks, source_coords, sink_coords = _get_sources_and_sinks(cg, data)
+        ret = cg.remove_edges(
+            user_id=user_id,
+            source_ids=sources,
+            sink_ids=sinks,
+            source_coords=source_coords,
+            sink_coords=sink_coords,
             mincut=mincut,
         )
     except cg_exceptions.LockingError as e:
@@ -459,7 +499,7 @@ def handle_split(table_id):
     current_app.operation_id = ret.operation_id
     if ret.new_root_ids is None:
         raise cg_exceptions.InternalServerError(
-            "Could not split selected segment groups."
+            f"{ret.operation_id}: Could not split selected segment groups."
         )
 
     current_app.logger.debug(("after split:", ret.new_root_ids))

pychunkedgraph/graph/cutting.py

@@ -1,23 +1,18 @@
-import collections
 import fastremap
 import numpy as np
 import itertools
-import logging
 import time
 import graph_tool
 import graph_tool.flow
 
-from typing import Dict
 from typing import Tuple
-from typing import Optional
 from typing import Sequence
 from typing import Iterable
 
 from .utils import flatgraph
-from .utils import basetypes
 from .utils.generic import get_bounding_box
 from .edges import Edges
-from .exceptions import PreconditionError
+from .exceptions import PreconditionError, SupervoxelSplitRequiredError
 from .exceptions import PostconditionError
 
 DEBUG_MODE = False
@@ -116,6 +111,10 @@ def __init__(
             self.cross_chunk_edge_remapping,
         ) = merge_cross_chunk_edges_graph_tool(cg_edges, cg_affs)
 
+        # save this representative mapping for supervoxel splitting
+        # passed along with SupervoxelSplitRequiredError
+        self.sv_remapping = dict(complete_mapping)
+
         dt = time.time() - time_start
         if logger is not None:
             logger.debug("Cross edge merging: %.2fms" % (dt * 1000))
@@ -233,9 +232,10 @@ def _augment_mincut_capacity(self):
                 self.source_graph_ids,
             )
         except AssertionError:
-            raise PreconditionError(
+            raise SupervoxelSplitRequiredError(
                 "Paths between source or sink points irreparably overlap other labels from other side. "
-                "Check that labels are correct and consider spreading points out farther."
+                "Check that labels are correct and consider spreading points out farther.",
+                self.sv_remapping
             )
 
         paths_e_s_no, paths_e_y_no, do_check = flatgraph.remove_overlapping_edges(
@@ -581,11 +581,12 @@ def _sink_and_source_connectivity_sanity_check(self, cut_edge_set):
                 # but return a flag to return a message to the user
                 illegal_split = True
             else:
-                raise PreconditionError(
+                raise SupervoxelSplitRequiredError(
                     "Failed to find a cut that separated the sources from the sinks. "
                     "Please try another cut that partitions the sets cleanly if possible. "
                     "If there is a clear path between all the supervoxels in each set, "
-                    "that helps the mincut algorithm."
+                    "that helps the mincut algorithm.",
+                    self.sv_remapping
                 )
         except IsolatingCutException as e:
             if self.split_preview:

pychunkedgraph/graph/exceptions.py

@@ -3,21 +3,25 @@
 
 class ChunkedGraphError(Exception):
     """Base class for all exceptions raised by the ChunkedGraph"""
+
     pass
 
 
 class LockingError(ChunkedGraphError):
     """Raised when a Bigtable Lock could not be acquired"""
+
     pass
 
 
 class PreconditionError(ChunkedGraphError):
     """Raised when preconditions for Chunked Graph operations are not met"""
+
     pass
 
 
 class PostconditionError(ChunkedGraphError):
     """Raised when postconditions for Chunked Graph operations are not met"""
+
     pass
 
 
@@ -42,7 +46,7 @@ def __init__(self, message):
         self.message = message
 
     def __str__(self):
-        return f'[{self.status_code}]: {self.message}'
+        return f"[{self.status_code}]: {self.message}"
 
 
 class ClientError(ChunkedGraphAPIError):
@@ -51,21 +55,25 @@ class ClientError(ChunkedGraphAPIError):
 
 class BadRequest(ClientError):
     """Exception mapping a ``400 Bad Request`` response."""
+
     status_code = http_client.BAD_REQUEST
 
 
 class Unauthorized(ClientError):
     """Exception mapping a ``401 Unauthorized`` response."""
+
     status_code = http_client.UNAUTHORIZED
 
 
 class Forbidden(ClientError):
     """Exception mapping a ``403 Forbidden`` response."""
+
     status_code = http_client.FORBIDDEN
 
 
 class Conflict(ClientError):
     """Exception mapping a ``409 Conflict`` response."""
+
     status_code = http_client.CONFLICT
 
 
@@ -75,9 +83,29 @@ class ServerError(ChunkedGraphAPIError):
 
 class InternalServerError(ServerError):
     """Exception mapping a ``500 Internal Server Error`` response."""
+
     status_code = http_client.INTERNAL_SERVER_ERROR
 
 
 class GatewayTimeout(ServerError):
     """Exception mapping a ``504 Gateway Timeout`` response."""
+
     status_code = http_client.GATEWAY_TIMEOUT
+
+
+class SupervoxelSplitRequiredError(ChunkedGraphError):
+    """
+    Raised when supervoxel splitting is necessary.
+    Edit process should catch this error and retry after supervoxel has been split.
+    Saves remapping required for detecting which supervoxels need to be split.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        sv_remapping: dict,
+        operation_id: int | None = None,
+    ):
+        super().__init__(message)
+        self.sv_remapping = sv_remapping
+        self.operation_id = operation_id

pychunkedgraph/graph/operation.py

@@ -26,7 +26,7 @@
 from .utils import serializers
 from .cache import CacheService
 from .cutting import run_multicut
-from .exceptions import PreconditionError
+from .exceptions import PreconditionError, SupervoxelSplitRequiredError
 from .exceptions import PostconditionError
 from .utils.generic import get_bounding_box as get_bbox
 from ..logging.log_db import TimeIt
@@ -451,6 +451,10 @@ def execute(
                         new_root_ids=new_root_ids,
                         new_lvl2_ids=new_lvl2_ids,
                     )
+            except SupervoxelSplitRequiredError as err:
+                raise SupervoxelSplitRequiredError(
+                    str(err), err.sv_remapping, operation_id=lock.operation_id
+                ) from err
             except PreconditionError as err:
                 self.cg.cache = None
                 raise PreconditionError(err) from err
@@ -852,9 +856,10 @@ def __init__(
         self.path_augment = path_augment
         self.disallow_isolating_cut = disallow_isolating_cut
         if np.any(np.in1d(self.sink_ids, self.source_ids)):
-            raise PreconditionError(
+            raise SupervoxelSplitRequiredError(
                 "Supervoxels exist in both sink and source, "
-                "try placing the points further apart."
+                "try placing the points further apart.",
+                None,
             )
 
         ids = np.concatenate([self.source_ids, self.sink_ids])

pychunkedgraph/repair/fake_edges.py

@@ -9,9 +9,9 @@
 from os import environ
 from typing import Optional
 
-environ["BIGTABLE_PROJECT"] = "<>"
-environ["BIGTABLE_INSTANCE"] = "<>"
-environ["GOOGLE_APPLICATION_CREDENTIALS"] = "<path>"
+# environ["BIGTABLE_PROJECT"] = "<>"
+# environ["BIGTABLE_INSTANCE"] = "<>"
+# environ["GOOGLE_APPLICATION_CREDENTIALS"] = "<path>"
 
 from pychunkedgraph.graph import edits
 from pychunkedgraph.graph import ChunkedGraph