Added tests

skimage/graph/tests/test_rag.py

@@ -2,7 +2,7 @@ import numpy as np
 from skimage import graph
 from skimage._shared.version_requirements import is_installed
 from numpy.testing.decorators import skipif
-from skimage import segmentation
+from skimage import segmentation, io
 from numpy import testing
 
 
@@ -108,3 +108,52 @@ def test_rag_error():
     labels[5:, :] = 1
     testing.assert_raises(ValueError, graph.rag_mean_color, img, labels,
                           2, 'non existant mode')
+
+
+@skipif(not is_installed('networkx'))
+def test_rag_error():
+    img = np.zeros((10, 10, 3), dtype='uint8')
+    labels = np.zeros((10, 10), dtype='uint8')
+    labels[:5, :] = 0
+    labels[5:, :] = 1
+    testing.assert_raises(ValueError, graph.rag_mean_color, img, labels,
+                          2, 'non existant mode')
+
+def _weight_mean_color(graph, src, dst, n):
+    #print 'merging
+    diff = graph.node[dst]['mean color'] - graph.node[n]['mean color']
+    diff = np.linalg.norm(diff)
+    return diff
+
+
+def _pre_merge_mean_color(graph, src, dst):
+    graph.node[dst]['total color'] += graph.node[src]['total color']
+    graph.node[dst]['pixel count'] += graph.node[src]['pixel count']
+    graph.node[dst]['mean color'] = (graph.node[dst]['total color'] /
+                                     graph.node[dst]['pixel count'])
+
+
+def merge_hierarchical_mean_color(labels, rag, thresh, rag_copy=True,
+                                  in_place_merge=False):
+    return graph.merge_hierarchical(labels, rag, thresh, rag_copy,
+                                    in_place_merge, _pre_merge_mean_color,
+                                    _weight_mean_color)
+
+@skipif(not is_installed('networkx'))
+def test_rag_hierarchical():
+    img = np.zeros((8, 8, 3), dtype='uint8')
+    labels = np.zeros((8, 8), dtype='uint8')
+
+    img[:, :, :] = 128
+    labels[:,:] = 1
+
+    img[0:4,0:4,:] = 255,255,255
+    labels[0:4, 0:4] = 2
+
+    img[4:, 0:4,:] = 0,0,0
+    labels[4:, 0:4] = 3
+
+    g = graph.rag_mean_color(img, labels)
+    result = merge_hierarchical_mean_color(labels, g, 300)
+    assert len(np.unique(result)) == 1
+    io.imsave('/home/vighnesh/Desktop/test.png', img)