--- a/main/image/src/bwlabeln.cc	Mon Sep 03 20:33:01 2012 +0000
+++ b/main/image/src/bwlabeln.cc	Tue Sep 04 01:58:03 2012 +0000
@@ -19,9 +19,6 @@
 #include <set>
 #include "union-find.h++"
 
-
-dim_vector conn26_dim(3,3,3);
-boolNDArray conn26(conn26_dim,1);
 typedef Array<octave_idx_type> coord;
 
 bool operator== (const coord& a, const coord& b)
@@ -83,8 +80,6 @@
   }
 };
 
-namespace {
-
 // A few basic utility functions
 //{
 inline
@@ -135,7 +130,7 @@
 operator- (const coord& a)
 {
   octave_idx_type na = a.nelem ();
-  coord retval( dim_vector(na,1) );
+  coord retval (dim_vector(na,1) );
   for (octave_idx_type i = 0; i < na; i++)
     {
       retval(i) = -a(i);
@@ -144,52 +139,10 @@
 }
 //}
 
-bool any_bad_argument (const octave_value_list& args)
-{
-  return false;
-
-  const int nargin = args.length ();
-  if (nargin < 1 || nargin > 2)
-    {
-      print_usage ();
-      return true;
-    }
-
-  if (!args (0).is_bool_type ())
-    {
-      error ("bwlabeln: first input argument must be a 'logical' ND-array");
-      return true;
-    }
-
-  if (nargin == 2)
-    {
-      if (!args (1).is_real_scalar () && ! args(1).is_bool_type())
-        {
-          error ("bwlabeln: second input argument must be a real scalar "
-                 "or a 'logical' connectivity array");
-            return true;
-        }
-    }
-
-  return false;
-}
-
-//debug
-#include <iostream>
-using namespace std;
-
-ostream&
-operator<< (ostream& os, const coord& aidx)
-{
-  for (octave_idx_type i = 0; i < aidx.nelem (); i++)
-    os << aidx(i) + 1 << " ";
-  return os;
-}
-
-set<coord>
+std::set<coord>
 populate_neighbours(const boolNDArray& conn_mask)
 {
-  set<coord> neighbours;
+  std::set<coord> neighbours;
 
   dim_vector conn_size = conn_mask.dims ();
   coord centre(dim_vector(conn_size.length (), 1), 1);
@@ -210,9 +163,114 @@
   return neighbours;
 }
 
+boolNDArray
+get_mask(int N){
+  bool* mask_ptr;
+  octave_idx_type n;
+
+  static bool mask4[] = {0, 1, 0,
+                         1, 0, 1,
+                         0, 1, 0};
+
+  static bool mask8[] = {1, 1, 1,
+                         1, 0, 1,
+                         1, 0, 1};
+
+  static bool mask6[] = {0, 0, 0,
+                         0, 1, 0,
+                         0, 0, 0,
+
+                         0, 1, 0,
+                         1, 0, 1,
+                         0, 1, 0,
+
+                         0, 0, 0,
+                         0, 1, 0,
+                         0, 0, 0};
+
+  static bool mask18[] = {0, 1, 0,
+                          1, 1, 1,
+                          0, 1, 0,
+
+                          1, 1, 1,
+                          1, 0, 1,
+                          1, 1, 1,
+
+                          0, 1, 0,
+                          1, 1, 1,
+                          0, 1, 0};
+
+  static bool mask26[] = {1, 1, 1,
+                          1, 1, 1,
+                          1, 1, 1,
+
+                          1, 1, 1,
+                          1, 0, 1,
+                          1, 1, 1,
+
+                          1, 1, 1,
+                          1, 1, 1,
+                          1, 1, 1};
+
+  switch (N){
+  case 4:
+    n = 2;
+    mask_ptr = mask4;
+    break;
+  case 8:
+    n = 2;
+    mask_ptr = mask8;
+    break;
+  case 6:
+    n = 3;
+    mask_ptr = mask6;
+    break;
+  case 18:
+    n = 3;
+    mask_ptr = mask18;
+    break;
+  case 26:
+    n = 3;
+    mask_ptr = mask26;
+    break;
+  default:
+    panic_impossible ();
+  }
+
+  boolNDArray conn_mask;
+  if (n == 2)
+    {
+      conn_mask.resize (dim_vector (3, 3));
+      for (octave_idx_type i = 0; i < 9; i++)
+        conn_mask(i) = mask_ptr[i];
+
+    }
+  else
+    {
+      conn_mask.resize (dim_vector (3, 3, 3));
+      for (octave_idx_type i = 0; i < 27; i++)
+        conn_mask(i) = mask_ptr[i];
+    }
+
+  return conn_mask;
+}
+
+boolNDArray
+get_mask (const boolNDArray& BW)
+{
+  dim_vector mask_dims = BW.dims();
+  for (auto i = 0; i < mask_dims.length (); i++)
+    mask_dims(i) = 3;
+
+  return boolNDArray (mask_dims, 1);
+}
+
+#include <iostream>
+
 DEFUN_DLD(bwlabeln, args, , "\
 -*- texinfo -*-\n\
-@deftypefn {Loadable Function} {[@var{l}, @var{num}] =} bwlabeln(@var{bw}, @var{n})\n\
+@deftypefn {Loadable Function} {[@var{l}, @var{num}] =} bwlabeln(@var{bw})\n\
+@deftypefnx {Loadable Function} {[@var{l}, @var{num}] =} bwlabeln(@var{bw}, @var{n})\n\
 Label foreground objects in the n-dimensional binary image @var{bw}.\n\
 \n\
 The optional argument @var{n} sets the connectivity and defaults 26.\n\
@@ -230,27 +288,78 @@
 
   union_find<coord, coord_hash> u_f;
 
-  if (any_bad_argument (args))
+  octave_idx_type nargin = args.length ();
+
+  if (nargin < 1 || nargin > 2)
+  {
+    print_usage ();
     return rval;
+  }
 
-  boolNDArray BW = args (0).bool_array_value ();
+  if (!args(0).is_bool_type ())
+    {
+      error ("bwlabeln: first input argument must be a 'logical' ND-array");
+      return rval;
+    }
 
+  boolNDArray BW = args(0).bool_array_value ();
   dim_vector size_vec = BW.dims ();
 
-  int nargin = args.length ();
-
   //Connectivity mask
   boolNDArray conn_mask;
-  if( nargin == 1)
-    conn_mask = conn26; //Implement this properly later
+
+  if (nargin == 2)
+    {
+      if (args(1).is_real_scalar ())
+        {
+          double N = args(1).scalar_value ();
+          if (size_vec.length () == 2 && N != 4 && N != 8)
+            error ("bwlabeln: for 2d arrays, scalar N must be 4 or 8");
+          else if (size_vec.length () == 3 && N != 6 && N != 18 && N != 26)
+            error ("bwlabeln: for 3d arrays, scalar N must be 4 or 8");
+          else if (size_vec.length () > 3)
+            error ("bwlabeln: for higher-dimensional arrays, N must be a "
+                   "connectivity mask");
+          else
+            conn_mask = get_mask (N);
+        }
+      else if (args(2).is_bool_type() )
+        {
+          conn_mask = args(2).bool_array_value ();
+          dim_vector conn_mask_dims = conn_mask.dims ();
+          if (conn_mask_dims.length () != size_vec.length ())
+            error ("bwlabeln: connectivity mask N must have the same "
+                   "dimensions as BW");
+          for (octave_idx_type i = 0; i < conn_mask_dims.length (); i++)
+            {
+              if (conn_mask_dims(i) != 3)
+                {
+                  error ("bwlabeln: connectivity mask N must have all "
+                         "dimensions equal to 3");
+                }
+            }
+        }
+      else
+        error ("bwlabeln: second input argument must be a real scalar "
+               "or a 'logical' connectivity array");
+    }
   else
-    conn_mask = conn26;
+    // Get the maximal mask that has same number of dims as BW.
+    conn_mask = get_mask (BW);
+
+  if (error_state)
+    return rval;
+
+  using std::set;
 
   set<coord> neighbours = populate_neighbours(conn_mask);
 
+  std::cout << "Union-finding..." << std::endl;;
+
+  bool* BW_vec = BW.fortran_vec ();
   for (octave_idx_type idx = 0; idx < BW.nelem (); idx++)
     {
-      if (BW(idx))
+      if (BW_vec[idx])
         {
           coord aidx = to_coord (size_vec, idx);
 
@@ -259,7 +368,7 @@
 
           //Replace this with C++0x range-based for loop later
           //(implemented in gcc 4.6)
-          for (auto nbr = neighbours.begin (); nbr!=  neighbours.end (); nbr++)
+          for (auto nbr = neighbours.begin (); nbr != neighbours.end (); nbr++)
             {
               coord n = *nbr + aidx;
               if (in_range (n,size_vec) and BW(n) )
@@ -297,4 +406,3 @@
   rval(1) = ids_to_label.size ();
   return rval;
 }
-}//anonymous namespace