From 34fc126b7aba0373824ad503b12a732f633e931d Mon Sep 17 00:00:00 2001
From: Constantin Pape <c.pape@gmx.net>
Date: Tue, 9 Jan 2024 21:08:53 +0100
Subject: [PATCH] Update lmc.py

---
 scripts/utils/lmc.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/scripts/utils/lmc.py b/scripts/utils/lmc.py
index 44387bb..26bbef9 100644
--- a/scripts/utils/lmc.py
+++ b/scripts/utils/lmc.py
@@ -3,6 +3,25 @@
 import elf.segmentation.features as feats
 import elf.segmentation.watershed as ws
 
+
+# if the compute_affinity features function doesn't work because of windows issues
+def affinity_feature_fallback(rag, ws, affs, offsets):
+    assert offsets == [[-1, 0, 0], [0, -1, 0], [0, 0, -1]]
+    xy_boundaries = np.maximum(affs[1], affs[2])
+    z_boundaries = affs[0]
+
+    # compute in plane features
+    features = feats.compute_boundary_mean_and_length(rag, xy_boundaries)
+    # compute between plane features
+    z_features = feats.compute_boundary_mean_and_length(rag, z_boundaries)
+
+    # over-ride between plane features in features
+    z_edges = feats.compute_z_edge_mask(rag, ws)
+    features[between_plane_edges] = z_features[z_edges]
+    
+    return features
+
+
 def mc_baseline(affs, fragments=None):
     affs = 1 - affs
     boundary_input = np.maximum(affs[1], affs[2])
@@ -16,6 +35,8 @@ def mc_baseline(affs, fragments=None):
             fragments[z] = wsz
     rag = feats.compute_rag(fragments)
     offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1]]
+    # This is only necessary if compute_affinity_features does not work due to issues on windows or mac.
+    # costs = affinity_feature_fallback(rag, fragments, affs, offsets)[:, 0]
     costs = feats.compute_affinity_features(rag, affs, offsets)[:, 0]
     edge_sizes = feats.compute_boundary_mean_and_length(rag, boundary_input)[:, 1]
     costs = mc.transform_probabilities_to_costs(costs, edge_sizes=edge_sizes)
@@ -36,7 +57,11 @@ def multicut_multi(affs, offsets=[[-1, 0, 0], [0, -1, 0], [0, 0, -1]], fragments
             offset += max_id
             fragments[z] = wsz
     rag = feats.compute_rag(fragments)
+
+    # This is only necessary if compute_affinity_features does not work due to issues on windows or mac.
+    # costs = affinity_feature_fallback(rag, fragments, affs, offsets)[:, 0]
     costs = feats.compute_affinity_features(rag, affs, offsets)[:, 0]
+    
     edge_sizes = feats.compute_boundary_mean_and_length(rag, boundary_input)[:, 1]
     costs = mc.transform_probabilities_to_costs(costs, edge_sizes=edge_sizes)
     node_labels = mc.multicut_kernighan_lin(rag, costs)