From 6314074f886167c9484e71b5fdad26a4ddb41616 Mon Sep 17 00:00:00 2001
From: Philipp Schlegel <observing@web.de>
Date: Tue, 17 Sep 2024 17:18:49 +0100
Subject: [PATCH] accept radius='auto' and map colors from skeleton to mesh

---
 octarine_navis_plugin/objects.py | 26 ++++++++++++++++++--------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/octarine_navis_plugin/objects.py b/octarine_navis_plugin/objects.py
index 4c7ce00..02ba4af 100644
--- a/octarine_navis_plugin/objects.py
+++ b/octarine_navis_plugin/objects.py
@@ -152,12 +152,20 @@ def neuron2gfx(x, color=None, random_ids=False, **kwargs):
         else:
             object_id = neuron.id  # this may also be a random ID
 
-        if kwargs.get("radius", False):
-            # Convert and carry connectors with us
-            if isinstance(neuron, navis.TreeNeuron):
-                _neuron = navis.conversion.tree2meshneuron(neuron)
-                _neuron.connectors = neuron.connectors
-                neuron = _neuron
+        if isinstance(neuron, navis.TreeNeuron) and kwargs.get('radius', False) == "auto":
+            # Number of nodes with radii
+            n_radii = (neuron.nodes.get("radius", pd.Series([])).fillna(0) > 0).sum()
+            # If less than 30% of nodes have a radius, we will fall back to lines
+            if n_radii / neuron.nodes.shape[0] < 0.3:
+                kwargs['radius'] = False
+
+            _neuron = navis.conversion.tree2meshneuron(neuron, warn_missing_radii=False)
+            _neuron.connectors = neuron.connectors
+            neuron = _neuron
+
+            # See if we need to map colors to vertices
+            if isinstance(colormap[i], np.ndarray) and colormap[i].ndim == 2:
+                colormap[i] = colormap[i][neuron.vertex_map]
 
         neuron_color = colormap[i]
         if not kwargs.get("connectors_only", False):
@@ -316,7 +324,9 @@ def skeleton2gfx(neuron, neuron_color, object_id, **kwargs):
         if neuron_color.ndim == 1:
             coords = navis.plotting.plot_utils.segments_to_coords(neuron)
         else:
-            coords, neuron_color = navis.plotting.plot_utils.segments_to_coords(neuron, node_colors=neuron_color)
+            coords, neuron_color = navis.plotting.plot_utils.segments_to_coords(
+                neuron, node_colors=neuron_color
+            )
             # `neuron_color` is now a list of colors for each segment; we have to flatten it
             # and add `None` to match the breaks
             neuron_color = np.vstack(
@@ -428,4 +438,4 @@ def skeletor2gfx(s, **kwargs):
     import navis
 
     s = navis.TreeNeuron(s, soma=None, id=getattr(s, "id", uuid.uuid4()))
-    return neuron2gfx(s, **kwargs)
\ No newline at end of file
+    return neuron2gfx(s, **kwargs)