uint32 support

thorbenk · Mar 22, 2014 · 83207e5 · 83207e5
1 parent 903995e
commit 83207e5
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 17 deletions.
diff --git a/blockedarray/blockwisecc_py.cxx b/blockedarray/blockwisecc_py.cxx
@@ -50,30 +50,28 @@
 using namespace BW;
 
 
-/* CC conversion */
-template <int N>
-void exportCCForDim()
+template <int N, class T>
+void exportSpecificCC(std::string suffix)
 {
     using namespace boost::python;
     using namespace BW;
-    typedef ConnectedComponents<N> BCC;
+    typedef ConnectedComponents<N, T> BCC;
 
-    class_<ConnectedComponents<N, vigra::UInt8> >("ConnectedComponents",
-                init<Source<N, vigra::UInt8>*, typename BCC::V>())
+    class_<BCC>(("ConnectedComponents"+suffix).c_str(),
+                init<Source<N, T>*, typename BCC::V>())
     .def("writeResult", &BCC::writeResult,
         (arg("hdf5file"), arg("hdf5group"), arg("compression")=1))
     .def("writeToSink", &BCC::writeToSink,
         (arg("sink")))
     ;
-    /*
-    class_<ConnectedComponents<N, vigra::UInt32> >("ConnectedComponents",
-                init<Source<N, vigra::UInt32>*, typename BCC::V>())
-    .def("writeResult", &BCC::writeResult,
-        (arg("hdf5file"), arg("hdf5group"), arg("compression")=1))
-    .def("writeToSink", &BCC::writeToSink,
-        (arg("sink")))
-    ;
-    */
+}
+
+/* CC conversion */
+template <int N>
+void exportCCForDim()
+{
+    exportSpecificCC<N, vigra::UInt8>("U8");
+    exportSpecificCC<N, vigra::UInt32>("U32");
 }
 
 /* ROI conversion */

diff --git a/blockedarray/opBlockedConnectedComponents.py b/blockedarray/opBlockedConnectedComponents.py
@@ -19,6 +19,9 @@ def __init__(self, *args, **kwargs):
         self._sourceMap3 = dict(zip(self.supportedDtypes, 
                                     [dim3.PySourceU8, dim3.PySourceU32]))
         self._sinkMap3 = {np.uint32: dim3.PySinkU32}
+        self._ccMap3 = dict(zip(self.supportedDtypes, 
+                                [dim3.ConnectedComponentsU8,
+                                 dim3.ConnectedComponentsU32]))
 
     def _updateSlice(self, c, t, bg):
         bg = self.Background[0, 0, 0, c, t].wait()
@@ -27,8 +30,8 @@ def _updateSlice(self, c, t, bg):
         sink = self._getSink(c, t)
         #TODO enable 2d
         blockShape = tuple([int(s) for s in self._cache.BlockShape.value[:3]])
-        print(blockShape)
-        cc = dim3.ConnectedComponents(source, blockShape)
+        CC = self._ccMap3[self.Input.meta.dtype]
+        cc = CC(source, blockShape)
 
         cc.writeToSink(sink)
 

diff --git a/blockedarray/testOpBlockedConnectedComponents.py b/blockedarray/testOpBlockedConnectedComponents.py
@@ -51,6 +51,24 @@ def testCorrectLabeling(self):
 
         assert_array_equal(vol, out)
 
+    def testCorrectLabelingInt(self):
+        vol = np.zeros((1000, 100, 10))
+        vol = vol.astype(np.uint32)
+        vol = vigra.taggedView(vol, axistags='xyz')
+
+        vol[20:40, 10:30, 2:4] = 1
+        vol = vol.withAxes(*'xyzct')
+
+        op = OpBlockedConnectedComponents(graph=Graph())
+        op.Input.setValue(vol)
+        op.Background.setValue(self.bg)
+
+        out = op.Output[...].wait()
+        tags = op.Output.meta.getTaggedShape()
+        out = vigra.taggedView(out, axistags=op.Output.meta.axistags)
+
+        assert_array_equal(vol, out)
+
 
 class TestSimpleThings(unittest.TestCase):
     def testRoi(self):