Better conversion from probability map to affinity

seung-lab · Nov 30, 2023 · a81bb3e · a81bb3e
1 parent 976cdf7
commit a81bb3e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/scripts/cut_chunk_common.py b/scripts/cut_chunk_common.py
@@ -42,7 +42,9 @@ def cut_data(data, start_coord, end_coord, padding):
     bb = tuple(slice(start_coord[i], end_coord[i]) for i in range(3))
     if data.shape[3] == 1:
         if data.dtype == 'float32':
-            return pad_data(numpy.stack([numpy.squeeze(data[bb])]*3, axis=-1), padding)
+            pmap = numpy.squeeze(data[bb])
+            affinity = [numpy.minimum(numpy.roll(pmap, shift=1, axis=axis), pmap) for axis in range(3)]
+            return pad_data(numpy.stack(affinity, axis=-1), padding)
         else:
             return pad_data(data[bb], padding)
     elif data.shape[3] == 3: