FindAdjacencies: Added option: 'nudge-points-inward'

flyemflows/workflow/findadjacencies.py

@@ -149,6 +149,12 @@ class FindAdjacencies(Workflow):
                 "oneOf": [{"type": "integer"}, {"type": "null"}],
                 "default": None
             },
+            "nudge-points-apart": {
+                "description": "Optionally 'nudge' the points apart from eachother, away from the\n"
+                               "exact edge boundary and slightly towards the interior of each object.\n",
+                "type": "boolean",
+                "default": False
+            },
             "cc-distance-threshold": {
                 "description": "When computing the connected components, don't use edges that exceed this distance threshold.\n"
                                "A threshold of 1.0 indicates that only direct adjacencies should be used.\n"
@@ -241,6 +247,7 @@ def execute(self):
         resource_config = self.config["resource-manager"]
         subset_requirement = options["subset-labels-requirement"]
         find_closest_using_scale = options["find-closest-using-scale"]
+        nudge = options["nudge-points-apart"]
 
         self.resource_mgr_client = ResourceManagerClient(resource_config["server"], resource_config["port"])
         volume_service = VolumeService.create_from_config(input_config, self.resource_mgr_client)
@@ -256,7 +263,7 @@ def find_adj(brick, sg):
                 #        it might be better to distribute each brick's rows.
                 #        (But that will involve a dask join step...)
                 #        The subset_groups should generally be under 1 GB, anyway...
-                edges = find_edges_in_brick(brick, None, sg, subset_requirement)
+                edges = find_edges_in_brick(brick, None, sg, subset_requirement, nudge)
                 brick.compress()
                 return edges
             adjacent_edge_tables = brickwall.bricks.map(find_adj, sg=delayed(subset_groups)).compute()
@@ -285,7 +292,7 @@ def find_adj(brick, sg):
         if find_closest_using_scale is not None:
             with Timer("Finding closest approaches", logger):
                 def find_closest(brick, sg):
-                    edges = find_edges_in_brick(brick, find_closest_using_scale, sg, subset_requirement)
+                    edges = find_edges_in_brick(brick, find_closest_using_scale, sg, subset_requirement, nudge)
                     brick.compress()
                     return edges
                 nonadjacent_edge_tables = brickwall.bricks.map(find_closest, sg=delayed(subset_groups)).compute()
@@ -453,7 +460,7 @@ def append_group_ccs(edges_df, subset_groups, max_distance=None):
         return edges_df, subset_groups
 
 
-def find_edges_in_brick(brick, closest_scale=None, subset_groups=[], subset_requirement=2):
+def find_edges_in_brick(brick, closest_scale=None, subset_groups=[], subset_requirement=2, nudge=False):
     """
     Find all pairs of adjacent labels in the given brick,
     and find the central-most point along the edge between them.
@@ -546,10 +553,37 @@ def find_edges_in_brick(brick, closest_scale=None, subset_groups=[], subset_requ
         np.save(f'problematic-remapped-brick-{brick_name}.npy', remapped_volume)
         logger.error(f"Error in brick (XYZ): {brick_name}") # This will appear in the worker log.
         raise
-    
+
     if best_edges_df is None:
         return None
 
+    if nudge:
+        # sign of direction from a -> b
+        s = np.sign(
+              best_edges_df[['zb', 'yb', 'xb']].values.astype(np.int32)
+            - best_edges_df[['za', 'ya', 'xa']].values.astype(np.int32)
+        )
+
+        # Nudge by a single voxel away from edge.
+        nudged_a = best_edges_df[['za', 'ya', 'xa']].values.astype(np.int32) - s
+        nudged_b = best_edges_df[['zb', 'yb', 'xb']].values.astype(np.int32) + s
+
+        nudged_a = np.maximum(nudged_a, (0,0,0))
+        nudged_a = np.minimum(nudged_a, np.array(remapped_volume.shape) - 1)
+
+        nudged_b = np.maximum(nudged_b, (0,0,0))
+        nudged_b = np.minimum(nudged_b, np.array(remapped_volume.shape) - 1)
+
+        # Keep only the nudged coordinates which retain the correct label;
+        # revert the others to the original coordinate.
+        keep_a = (remapped_volume[(*nudged_a.T,)] == best_edges_df['label_a'])
+        keep_b = (remapped_volume[(*nudged_b.T,)] == best_edges_df['label_b'])
+
+        best_edges_df.loc[keep_a, ['za', 'ya', 'xa']] = nudged_a[keep_a.values]
+        best_edges_df.loc[keep_b, ['zb', 'yb', 'xb']] = nudged_b[keep_b.values]
+
+
+
     # Translate coordinates to global space
     best_edges_df.loc[:, ['za', 'ya', 'xa']] += brick.physical_box[0]
     best_edges_df.loc[:, ['zb', 'yb', 'xb']] += brick.physical_box[0]
@@ -560,7 +594,7 @@ def find_edges_in_brick(brick, closest_scale=None, subset_groups=[], subset_requ
 
     # Normalize
     swap_df_cols(best_edges_df, None, best_edges_df.eval('label_a > label_b'), ['a', 'b'])
-    
+
     return best_edges_df
 
 

tests/workflows/test_findadjacencies.py

@@ -169,6 +169,37 @@ def test_findadjacencies_subset_edges(setup_findadjacencies):
     assert (output_df.query('label_a == 3')[['xa', 'xb']].values[0] == (2,2)).all()
 
 
+def test_findadjacencies_subset_edges_and_nudge(setup_findadjacencies):
+    template_dir, config, _volume = setup_findadjacencies
+
+    subset_edges = pd.DataFrame([[4,3]], columns=['label_a', 'label_b'])
+    subset_edges.to_csv(f'{template_dir}/subset-edges.csv', index=False, header=True)
+
+    # Overwrite config with updated settings.
+    config = copy.deepcopy(config)
+    config["findadjacencies"]["subset-edges"] = 'subset-edges.csv'
+    config["findadjacencies"]["nudge-points-apart"] = True
+
+    with open(f"{template_dir}/workflow.yaml", 'w') as f:
+        yaml.dump(config, f)
+
+    execution_dir, workflow = launch_flow(template_dir, 1)
+
+    final_config = workflow.config
+    output_df = pd.read_csv(f'{execution_dir}/{final_config["findadjacencies"]["output-table"]}')
+
+    label_pairs = output_df[['label_a', 'label_b']].values
+    assert 0 not in label_pairs.flat
+
+    label_pairs = list(map(tuple, label_pairs))
+    assert (1,2) not in label_pairs
+    assert (3,4) in label_pairs
+
+    assert (output_df.query('label_a == 3')[['za', 'zb']].values[0] == (0,0)).all()
+    assert (output_df.query('label_a == 3')[['ya', 'yb']].values[0] == (6,5)).all() # not 'forward'
+    assert (output_df.query('label_a == 3')[['xa', 'xb']].values[0] == (2,2)).all()
+
+
 def test_findadjacencies_solid_volume():
     """
     If the volume is solid or empty, an error is raised at the end of the workflow.
@@ -499,7 +530,7 @@ def test_append_group_col():
     CLUSTER_TYPE = os.environ['CLUSTER_TYPE'] = "synchronous"
     args = ['-s', '--tb=native', '--pyargs', 'tests.workflows.test_findadjacencies']
     #args += ['-x']
-    #args += ['-k append_group_col']
+    #args += ['-k test_findadjacencies_subset_edges_and_nudge']
     #args += ['-k', 'findadjacencies_from_dvid_sparse_groups'
     #         ' or findadjacencies_different_dvid_blocks_sparse_labels'
     #         ' or findadjacencies_from_dvid_sparse_edges'