--- a/src/__bwdist.cc
+++ b/src/__bwdist.cc
@@ -1,163 +1,192 @@
+// bwdist.cc - OCT file, implements the BWDIST function
+// Depends on "edtfunc.c" for the actual computations
 /*
-Copyright (C) 2006 Pedro Felzenszwalb
+
+Copyright (C) 2009 Stefan Gustavson (stefan.gustavson@gmail.com)
 
-This program is free software; you can redistribute it and/or modify
-it under the terms of the GNU General Public License as published by
-the Free Software Foundation; either version 2 of the License, or
-(at your option) any later version.
+This program is free software; you can redistribute it and/or modify it
+under the terms of the GNU General Public License as published by the
+Free Software Foundation; either version 3 of the License, or (at your
+option) any later version.
 
-This program is distributed in the hope that it will be useful,
-but WITHOUT ANY WARRANTY; without even the implied warranty of
-MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-GNU General Public License for more details.
+This program is distributed in the hope that it will be useful, but WITHOUT
+ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+for more details.
 
 You should have received a copy of the GNU General Public License
-along with this program; If not, see <http://www.gnu.org/licenses/>.
-*/
+along with Octave; see the file COPYING.  If not, see
+<http://www.gnu.org/licenses/>.
 
-/* 
- * Remark by Søren Hauberg, december 28th 2006.
- *
- * This code was mainly written by Pedro Felzenszwalb and published
- * on http://people.cs.uchicago.edu/~pff/dt/
- * Pedro kindly released his code under the GPL and I ported it to
- * Octave.
- */
+*/
 
 #include <octave/oct.h>
 
-#define INF 1E20
 
-template <class T>
-inline T square(const T &x) { return x*x; };
-
-/* dt of 1d function using squared distance */
-static float* dt (float *f, int n)
+#ifdef __cplusplus
+extern "C"
 {
-  float *d = new float[n];
-  int *v = new int[n];
-  float *z = new float[n+1];
-  int k = 0;
-  v[0] = 0;
-  z[0] = -INF;
-  z[1] = +INF;
-  for (int q = 1; q <= n-1; q++)
-    {
-      float s  = ((f[q] + square (q)) - (f[v[k]] + square (v[k])))/(2*q-2*v[k]);
-      while (s <= z[k])
-        {
-          k--;
-          s  = ((f[q] + square (q)) - (f[v[k]] + square (v[k])))/(2*q-2*v[k]);
-        }
-      k++;
-      v[k] = q;
-      z[k] = s;
-      z[k+1] = +INF;
-    }
+#endif
+
+#define DIST_EUCLIDEAN(x,y) ((int)(x)*(x) + (y)*(y))
+#define MAX(x,y) ((x)>(y) ? (x) : (y))
+#define DIST_CHESSBOARD(x,y) (MAX(abs(x), abs(y)))
+#define DIST_CITYBLOCK(x,y) (abs(x) + abs(y))
+#define SQRT2_1 0.4142136536
+#define DIST_QUASI_EUCLIDEAN(x,y) (abs(x)>abs(y) ? (abs(x) + SQRT2_1 * abs(y)) : (SQRT2_1 * abs(x) + abs(y)))
+
+#define DIST(x,y) DIST_EUCLIDEAN(x,y)
+#define FUNCNAME euclidean
+#include "edtfunc.c"
+#undef DIST
+#undef FUNCNAME
+
+#define DIST(x,y) DIST_CHESSBOARD(x,y)
+#define FUNCNAME chessboard
+#include "edtfunc.c"
+#undef DIST
+#undef FUNCNAME
+
+#define DIST(x,y) DIST_CITYBLOCK(x,y)
+#define FUNCNAME cityblock
+#include "edtfunc.c"
+#undef DIST
+#undef FUNCNAME
+
+#define DIST(x,y) DIST_QUASI_EUCLIDEAN(x,y)
+#define FUNCNAME quasi_euclidean
+#include "edtfunc.c"
+#undef DIST
+#undef FUNCNAME
+
+#ifdef __cplusplus
+}  /* end extern "C" */
+#endif
 
-  k = 0;
-  for (int q = 0; q <= n-1; q++)
-    {
-      while (z[k+1] < q)
-        k++;
-      d[q] = square (q-v[k]) + f[v[k]];
-    }
-
-  delete [] v;
-  delete [] z;
-  return d;
-}
-
-/* dt of 2d function using squared distance */
-static void dt (NDArray &im)
+DEFUN_DLD (bwdist, args, nargout,
+"-*- texinfo -*-\n\
+@deftypefn {Function File} {@var{D} =} bwdist(@var{bw})\n\
+Computes the distance transform of the image @var{bw}.\n\
+@var{bw} should be a binary 2D array, either a Boolean array or a\n\
+numeric array containing only the values 0 and 1.\n\
+The return value @var{D} is a double matrix of the same size as @var{bw}.\n\
+Elements with value 0 are considered background pixels, elements\n\
+with value 1 are considered object pixels. The return value\n\
+for each background pixel is the distance (according to the chosen\n\
+metric) to the closest object pixel. For each object pixel the\n\
+return value is 0.\n\
+\n\
+@deftypefnx{Function File} {@var{D} =} bwdist(@var{bw}, @var{method})\n\
+\n\
+@var{method} is a string to choose the distance metric. Currently\n\
+available metrics are 'euclidean', 'chessboard', 'cityblock' and\n\
+'quasi-euclidean', which may each be abbreviated\n\
+to any string starting with 'e', 'ch', 'ci' and 'q', respectively.\n\
+If @var{method} is not specified, 'euclidean' is the default.\n\
+\n\
+@deftypefnx {Function File} {[@var{D},@var{C}] =} bwdist(@var{bw}, @var{method})\n\
+\n\
+If a second output argument is given, the linear index for the\n\
+closest object pixel is returned for each pixel. (For object\n\
+pixels, the index points to the pixel itself.) The return value C\n\
+is a matrix the same size as @var{bw}.\n\n\
+@end deftypefn")
 {
-  const int width = im.dim1 ();
-  const int height = im.dim2 ();
-  float *f = new float[std::max(width,height)];
-
-  // transform along columns
-  for (int x = 0; x < width; x++)
-    {
-      for (int y = 0; y < height; y++)
-        {
-          f[y] = im (x, y);
-        }
-      float *d = dt (f, height);
-      for (int y = 0; y < height; y++)
-        {
-          im (x, y) = d[y];
-        }
-      delete [] d;
+  const int nargin = args.length();
+  octave_value_list retval;
+  
+  /* Check for proper number of input and output arguments */
+  if ((nargin < 1) || (nargin>2)) {
+    error ("bwdist accepts only one or two input parameters.");
+  }
+  else if (nargout > 2) {
+    error ("bwdist returns at most 2 output parameters.");
+  }
+  else {
+    /* Make sure input is a matrix */
+    const Matrix bw = args(0).matrix_value();
+    if (error_state) {
+      error ("bwdist input argument must be a matrix");
+      return retval;
+    }
+    /* Warn if input is not a binary image */
+    if(bw.any_element_not_one_or_zero()) {
+      warning ("bwdist input contains values other than 1 and 0.");
     }
 
-    // transform along rows
-    for (int y = 0; y < height; y++)
-      {
-        for (int x = 0; x < width; x++)
-          {
-            f[x] = im (x, y);
-          }
-        float *d = dt (f, width);
-        for (int x = 0; x < width; x++)
-          {
-            im (x, y) = d[x];
-          }
-        delete [] d;
+    /* Everything seems to be OK to proceed */
+    dim_vector dims = bw.dims();
+    int rows = dims(0);
+    int cols = dims(1);
+    int caseMethod = 0; // Default 0 means Euclidean
+    if(nargin > 1) {
+      charMatrix method = args(1).char_matrix_value();
+      if(method(0) == 'e') caseMethod = 0; // Euclidean;
+      else if (method(0) == 'c') {
+	if(method(1) == 'h') caseMethod = 1; // chessboard
+	else if(method(1) == 'i') caseMethod = 2; // cityblock
       }
-
-    delete [] f;
-}
-
-
-/* dt of binary image using squared distance */
-void dt (const boolNDArray &im, NDArray &out)
-{
-  const int width = im.dim1 ();
-  const int height = im.dim2 ();
-
-  for (int y = 0; y < height; y++)
-    {
-      for (int x = 0; x < width; x++)
-        {
-          if (im (x, y))
-            out (x, y) = 0;
-          else
-            out (x, y) = INF;
-        }
+      else if(method(0) == 'q') caseMethod = 3; // quasi-Euclidean
+      else {
+	warning ("unknown metric, using 'euclidean'");
+	caseMethod = 0;
+      }
     }
 
-  dt (out);
-}
+    if (!error_state) {
+      /* Allocate two arrays for temporary output values */
+      OCTAVE_LOCAL_BUFFER (short, xdist, dims.numel());
+      OCTAVE_LOCAL_BUFFER (short, ydist, dims.numel());
+
+      /* Create final output array */
+      Matrix D (rows, cols);
+
+      /* Call the appropriate C subroutine and compute output */
+      switch(caseMethod) {
+
+      case 1:
+	chessboard(bw, rows, cols, xdist, ydist);
+	for(int i=0; i<rows*cols; i++) {
+	  D(i) = DIST_CHESSBOARD(xdist[i], ydist[i]);
+	}
+	break;
+
+      case 2:
+	cityblock(bw, rows, cols, xdist, ydist);
+	for(int i=0; i<rows*cols; i++) {
+	  D(i) = DIST_CITYBLOCK(xdist[i], ydist[i]);
+	}
+	break;
 
-DEFUN_DLD (__bwdist, args, nargout, "\
--*- texinfo -*-\n\
-@deftypefn {Function File} __bwdist (@var{bw})\n\
-Computes the Euclidian Distance Transform for a binary image @var{bw}.\n\
-You should not call this function directly, instead call 'bwdist'.\n\
-@seealso{bwdist}\n\
-@end deftypefn\n\
-")
-{
-  octave_value_list retval;
-  
-  const boolNDArray bw = args(0).bool_array_value ();
-  if (error_state)
-    {
-      error ("__bwdist: input must be a boolean matrix");
-      return retval;
+      case 3:
+	quasi_euclidean(bw, rows, cols, xdist, ydist);
+	for(int i=0; i<rows*cols; i++) {
+	  D(i) = DIST_QUASI_EUCLIDEAN(xdist[i], ydist[i]);
+	}
+	break;
+
+      case 0:
+      default:
+	euclidean(bw, rows, cols, xdist, ydist);
+	/* Remember sqrt() for the final output */
+	for(int i=0; i<rows*cols; i++) {
+	  D(i) = sqrt((double)DIST_EUCLIDEAN(xdist[i], ydist[i]));
+	}
+	break;
+      }
+
+      retval(0) = D;
+
+      if(nargout > 1) {
+	/* Create a second output array */
+	Matrix C (rows, cols);
+	/* Compute optional 'index to closest object pixel' */
+	for(int i=0; i<rows*cols; i++) {
+	  C (i) = i+1 - xdist[i] - ydist[i]*rows;
+	}
+	retval(1) = C;
+      }
     }
-  
-  const dim_vector dims = bw.dims ();
-  if (dims.length () != 2)
-    {
-      error ("__bwdist: currently only binary images are supported");
-      return retval;
-    }
-  
-  NDArray out (dims);
-  
-  dt (bw, out);
-  
-  retval.append (out);
+  }
   return retval;
 }

new file mode 100755
--- /dev/null
+++ b/src/edtfunc.c
@@ -0,0 +1,380 @@
+// edtfunc.c - Euclidean distance transform of a binary image
+/*
+
+Copyright (C) 2009 Stefan Gustavson (stefan.gustavson@gmail.com)
+
+This program is free software; you can redistribute it and/or modify it
+under the terms of the GNU General Public License as published by the
+Free Software Foundation; either version 3 of the License, or (at your
+option) any later version.
+
+This program is distributed in the hope that it will be useful, but WITHOUT
+ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+for more details.
+
+You should have received a copy of the GNU General Public License
+along with Octave; see the file COPYING.  If not, see
+<http://www.gnu.org/licenses/>.
+
+*/
+
+/*
+ * This is a sweep-and-update Euclidean distance transform of
+ * a binary image. All positive pixels are considered object
+ * pixels, zero or negative pixels are treated as background.
+ *
+ * By Stefan Gustavson (stefan.gustavson@gmail.com).
+ *
+ * Originally written in 1994, based on paper-only descriptions
+ * of the SSED8 algorithm, invented by Per-Erik Danielsson
+ * and improved by Ingemar Ragnemalm. This is a classic algorithm
+ * with roots in the 1980s, still very good for the 2D case.
+ *
+ * Updated in 2004 to treat pixels at image edges correctly,
+ * and to improve code readability.
+ *
+ * Edited in 2009 to form the foundation for Octave BWDIST:
+ * added #define-configurable distance measure and function name
+ *
+ */
+
+#ifndef DIST
+#define DIST(x,y) ((int)(x) * (x) + (y) * (y))
+#endif
+
+#ifndef FUNCNAME
+#define FUNCNAME edt
+#endif
+
+void FUNCNAME(const Matrix &img, int w, int h, short *distx, short *disty)
+{
+  int x, y, i;
+  int offset_u, offset_ur, offset_r, offset_rd,
+  offset_d, offset_dl, offset_l, offset_lu;
+  double olddist2, newdist2, newdistx, newdisty;
+  int changed;
+
+  /* Initialize index offsets for the current image width */
+  offset_u = -w;
+  offset_ur = -w+1;
+  offset_r = 1;
+  offset_rd = w+1;
+  offset_d = w;
+  offset_dl = w-1;
+  offset_l = -1;
+  offset_lu = -w-1;
+
+  /* Initialize the distance images to be all large values */
+  for(i=0; i<w*h; i++)
+    if(img(i) <= 0.0)
+      {
+        distx[i] = 32000; // Large but still representable in a short, and
+        disty[i] = 32000; // 32000^2 + 32000^2 does not overflow an int
+      }
+    else
+      {
+        distx[i] = 0;
+        disty[i] = 0;
+      }
+
+  /* Perform the transformation */
+  do
+    {
+      changed = 0;
+
+      /* Scan rows, except first row */
+      for(y=1; y<h; y++)
+        {
+
+          /* move index to leftmost pixel of current row */
+          i = y*w; 
+
+          /* scan right, propagate distances from above & left */
+
+          /* Leftmost pixel is special, has no left neighbors */
+          olddist2 = DIST(distx[i], disty[i]);
+          if(olddist2 > 0) // If not already zero distance
+            {
+              newdistx = distx[i+offset_u];
+              newdisty = disty[i+offset_u]+1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_ur]-1;
+              newdisty = disty[i+offset_ur]+1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  changed = 1;
+                }
+            }
+          i++;
+
+          /* Middle pixels have all neighbors */
+          for(x=2; x<w-1; x++, i++)
+            {
+              olddist2 = DIST(distx[i], disty[i]);
+              if(olddist2 == 0) continue; // Already zero distance
+
+              newdistx = distx[i+offset_l]+1;
+              newdisty = disty[i+offset_l];
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_lu]+1;
+              newdisty = disty[i+offset_lu]+1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_u];
+              newdisty = disty[i+offset_u]+1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_ur]-1;
+              newdisty = disty[i+offset_ur]+1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  changed = 1;
+                }
+            }
+
+          /* Rightmost pixel of row is special, has no right neighbors */
+          i++;
+          olddist2 = DIST(distx[i], disty[i]);
+          if(olddist2 > 0) // If not already zero distance
+            {
+              newdistx = distx[i+offset_l]+1;
+              newdisty = disty[i+offset_l];
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_lu]+1;
+              newdisty = disty[i+offset_lu]+1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_u];
+              newdisty = disty[i+offset_u]+1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+            }
+
+          /* Move index to second rightmost pixel of current row. */
+          /* Rightmost pixel is skipped, it has no right neighbor. */
+          i = y*w + w-2;
+
+          /* scan left, propagate distance from right */
+          for(x=w-2; x>=0; x--, i--)
+            {
+              olddist2 = DIST(distx[i], disty[i]);
+              if(olddist2 == 0) continue; // Already zero distance
+              
+              newdistx = distx[i+offset_r]-1;
+              newdisty = disty[i+offset_r];
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  changed = 1;
+                }
+            }
+        }
+      
+      /* Scan rows in reverse order, except last row */
+      for(y=h-2; y>=0; y--)
+        {
+          /* move index to rightmost pixel of current row */
+          i = y*w + w-1;
+
+          /* Scan left, propagate distances from below & right */
+
+          /* Rightmost pixel is special, has no right neighbors */
+          olddist2 = DIST(distx[i], disty[i]);
+          if(olddist2 > 0) // If not already zero distance
+            {
+              newdistx = distx[i+offset_d];
+              newdisty = disty[i+offset_d]-1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_dl]+1;
+              newdisty = disty[i+offset_dl]-1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  changed = 1;
+                }
+            }
+          i--;
+
+          /* Middle pixels have all neighbors */
+          for(x=w-2; x>0; x--, i--)
+            {
+              olddist2 = DIST(distx[i], disty[i]);
+              if(olddist2 == 0) continue; // Already zero distance
+
+              newdistx = distx[i+offset_r]-1;
+              newdisty = disty[i+offset_r];
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_rd]-1;
+              newdisty = disty[i+offset_rd]-1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_d];
+              newdisty = disty[i+offset_d]-1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_dl]+1;
+              newdisty = disty[i+offset_dl]-1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  changed = 1;
+                }
+            }
+          /* Leftmost pixel is special, has no left neighbors */
+          olddist2 = DIST(distx[i], disty[i]);
+          if(olddist2 > 0) // If not already zero distance
+            {
+              newdistx = distx[i+offset_r]-1;
+              newdisty = disty[i+offset_r];
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_rd]-1;
+              newdisty = disty[i+offset_rd]-1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+
+              newdistx = distx[i+offset_d];
+              newdisty = disty[i+offset_d]-1;
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  olddist2=newdist2;
+                  changed = 1;
+                }
+            }
+
+          /* Move index to second leftmost pixel of current row. */
+          /* Leftmost pixel is skipped, it has no left neighbor. */
+          i = y*w + 1;
+          for(x=1; x<w; x++, i++)
+            {
+              /* scan right, propagate distance from left */
+              olddist2 = DIST(distx[i], disty[i]);
+              if(olddist2 == 0) continue; // Already zero distance
+
+              newdistx = distx[i+offset_l]+1;
+              newdisty = disty[i+offset_l];
+              newdist2 = DIST(newdistx, newdisty);
+              if(newdist2 < olddist2)
+                {
+                  distx[i]=newdistx;
+                  disty[i]=newdisty;
+                  changed = 1;
+                }
+            }
+        }
+    }
+  while(changed); // Sweep until no more updates are made
+
+  /* The transformation is completed. */
+
+}