Improve efficiency of colorize_labels tool wrapper + bugfix

tools/colorize_labels/colorize_labels.py

@@ -1,12 +1,11 @@
 import argparse
 
 import giatools.io
-import matplotlib.colors as mpl
+import matplotlib.pyplot as plt
 import networkx as nx
 import numpy as np
 import scipy.ndimage as ndi
 import skimage.io
-import skimage.morphology as morph
 import skimage.util
 
 
@@ -22,7 +21,6 @@ def color_hex_to_rgb_tuple(hex):
 
 def build_label_adjacency_graph(im, radius, bg_label):
     G = nx.Graph()
-    selem = morph.disk(radius)
     for label in np.unique(im):
 
         if label == bg_label:
@@ -31,7 +29,7 @@ def build_label_adjacency_graph(im, radius, bg_label):
         G.add_node(label)
 
         cc = (im == label)
-        neighborhood = ndi.binary_dilation(cc, selem)
+        neighborhood = (ndi.distance_transform_edt(~cc) <= radius)
         adjacent_labels = np.unique(im[neighborhood])
 
         for adjacent_label in adjacent_labels:
@@ -45,6 +43,18 @@ def build_label_adjacency_graph(im, radius, bg_label):
     return G
 
 
+def get_n_unique_mpl_colors(n, colormap='jet', cyclic=False):
+    """
+    Yields `n` unique colors from the given `colormap`.
+
+    Set `cyclic` to `True` if the `colormap` is cyclic.
+    """
+    cmap = plt.get_cmap(colormap)
+    m = n if cyclic else n - 1
+    for i in range(n):
+        yield np.multiply(255, cmap(i / m))
+
+
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser()
@@ -68,7 +78,7 @@ def build_label_adjacency_graph(im, radius, bg_label):
     unique_colors = frozenset(graph_coloring.values())
 
     # Assign colors to nodes based on the greedy coloring
-    graph_color_to_mpl_color = dict(zip(unique_colors, mpl.TABLEAU_COLORS.values()))
+    graph_color_to_mpl_color = dict(zip(unique_colors, get_n_unique_mpl_colors(len(unique_colors))))
     node_colors = [graph_color_to_mpl_color[graph_coloring[n]] for n in G.nodes()]
 
     # Render result
@@ -77,7 +87,6 @@ def build_label_adjacency_graph(im, radius, bg_label):
     for label, label_color in zip(G.nodes(), node_colors):
 
         cc = (im == label)
-        label_color = color_hex_to_rgb_tuple(label_color)
         for ch in range(3):
             result[:, :, ch][cc] = label_color[ch]
 

tools/colorize_labels/colorize_labels.xml

@@ -4,7 +4,7 @@
         <import>creators.xml</import>
         <import>tests.xml</import>
         <token name="@TOOL_VERSION@">3.2.1</token>
-        <token name="@VERSION_SUFFIX@">2</token>
+        <token name="@VERSION_SUFFIX@">3</token>
     </macros>
     <creator>
         <expand macro="creators/bmcv" />
@@ -43,6 +43,7 @@
        <data format="png" name="output" from_work_dir="output.png" />
     </outputs>
     <tests>
+        <!-- int64 -->
         <test>
             <param name="input" value="input1.tif" />
             <param name="radius" value="1" />
@@ -52,6 +53,7 @@
                 <has_image_channels channels="3"/>
             </expand>
         </test>
+        <!-- uint8 -->
         <test>
             <param name="input" value="input2.tif" />
             <param name="radius" value="10" />
@@ -61,6 +63,16 @@
                 <has_image_channels channels="3"/>
             </expand>
         </test>
+        <!-- uint16 -->
+        <test>
+            <param name="input" value="input3.tif" />
+            <param name="radius" value="100" />
+            <param name="bg_label" value="0" />
+            <param name="bg_color" value="#ffffff" />
+            <expand macro="tests/intensity_image_diff" name="output" value="output3.png" ftype="png">
+                <has_image_channels channels="3"/>
+            </expand>
+        </test>
     </tests>
     <help>