From 01145183a023f84fe9625c673428971a4f69e621 Mon Sep 17 00:00:00 2001
From: William Silversmith <william.silversmith@gmail.com>
Date: Thu, 16 Oct 2025 20:06:39 -0400
Subject: [PATCH 1/4] feat: multilabel support

---
 src/auxiliary.hpp |  5 +--
 src/fastxs3d.cpp  | 41 +++++++++++++++++-------
 src/xs3d.hpp      | 81 +++++++++++++++++++++++++----------------------
 xs3d/__init__.py  | 22 +++++++------
 4 files changed, 89 insertions(+), 60 deletions(-)

diff --git a/src/auxiliary.hpp b/src/auxiliary.hpp
index f720dc1..7b7f9f4 100644
--- a/src/auxiliary.hpp
+++ b/src/auxiliary.hpp
@@ -190,8 +190,9 @@ std::tuple<float*, uint8_t> cross_section_slow(
  * connected component, so pre-filtering must be performed to 
  * ensure a match.
  */
+template <typename LABEL>
 std::tuple<float, uint8_t> cross_sectional_area_slow(
-	const uint8_t* binimg,
+	const LABEL* labels, const LABEL segid,
 	const uint64_t sx, const uint64_t sy, const uint64_t sz,
 	
 	const float px, const float py, const float pz,
@@ -239,7 +240,7 @@ std::tuple<float, uint8_t> cross_sectional_area_slow(
 			for (uint64_t x = 0; x < sx; x++) {
 				uint64_t loc = x + sx * (y + sy * z);
 
-				if (!binimg[loc]) {
+				if (labels[loc] != segid) {
 					continue;
 				}
 
diff --git a/src/fastxs3d.cpp b/src/fastxs3d.cpp
index 486e589..a29b19c 100644
--- a/src/fastxs3d.cpp
+++ b/src/fastxs3d.cpp
@@ -79,21 +79,23 @@ auto section(
 	return std::make_tuple(arr, std::get<1>(tup));
 }
 
+template <typename LABEL>
 auto calculate_area(
-	const py::array_t<uint8_t> &binimg,
+	const py::array_t<LABEL> &image,
+	const LABEL segid,
 	const py::array_t<float> &point,
 	const py::array_t<float> &normal,
 	const py::array_t<float> &anisotropy,
 	const bool slow_method,
 	const bool use_persistent_data
 ) {
-	const uint64_t sx = binimg.shape()[0];
-	const uint64_t sy = binimg.ndim() < 2
+	const uint64_t sx = image.shape()[0];
+	const uint64_t sy = image.ndim() < 2
 		? 1 
-		: binimg.shape()[1];
-	const uint64_t sz = binimg.ndim() < 3 
+		: image.shape()[1];
+	const uint64_t sz = image.ndim() < 3 
 		? 1 
-		: binimg.shape()[2];
+		: image.shape()[2];
 
 	if (point.size() < 3) {
 	    throw py::value_error("point array must have at least 3 elements");
@@ -112,8 +114,8 @@ auto calculate_area(
 	}
 
 	if (slow_method) {
-		return xs3d::cross_sectional_area_slow(
-			binimg.data(),
+		return xs3d::cross_sectional_area_slow<LABEL>(
+			image.data(), segid,
 			sx, sy, sz,
 			point.at(0), point.at(1), point.at(2),
 			normal.at(0), normal.at(1), normal.at(2),
@@ -121,8 +123,8 @@ auto calculate_area(
 		);
 	}
 	else {
-		return xs3d::cross_sectional_area(
-			binimg.data(),
+		return xs3d::cross_sectional_area<LABEL>(
+			image.data(), segid,
 			sx, sy, sz,
 			point.at(0), point.at(1), point.at(2),
 			normal.at(0), normal.at(1), normal.at(2),
@@ -231,12 +233,29 @@ auto projection(
     }
 }
 
+
+#define REGISTER_CALCULATE_AREA_FOR_TYPE(T) \
+    m.def("area", &calculate_area<T>, "Find the cross sectional area for a given image, point, and normal vector.");
+
 PYBIND11_MODULE(fastxs3d, m) {
 	m.doc() = "Finding cross sectional area in 3D voxelized images."; 
 	m.def("projection", &projection, "Project a cross section of a 3D image onto a 2D plane");
 	m.def("section", &section, "Return a floating point image that shows the voxels contributing area to a cross section.");
-	m.def("area", &calculate_area, "Find the cross sectional area for a given binary image, point, and normal vector.");
+
+	REGISTER_CALCULATE_AREA_FOR_TYPE(bool)
+	REGISTER_CALCULATE_AREA_FOR_TYPE(uint8_t)
+	REGISTER_CALCULATE_AREA_FOR_TYPE(uint16_t)
+	REGISTER_CALCULATE_AREA_FOR_TYPE(uint32_t)
+	REGISTER_CALCULATE_AREA_FOR_TYPE(uint64_t)
+	REGISTER_CALCULATE_AREA_FOR_TYPE(int8_t)
+	REGISTER_CALCULATE_AREA_FOR_TYPE(int16_t)
+	REGISTER_CALCULATE_AREA_FOR_TYPE(int32_t)
+	REGISTER_CALCULATE_AREA_FOR_TYPE(int64_t)
+	REGISTER_CALCULATE_AREA_FOR_TYPE(float)
+	REGISTER_CALCULATE_AREA_FOR_TYPE(double)
 	
 	m.def("set_shape", &xs3d::set_shape, "Accelerate the area function across many evaluation points by saving some attributes of the input shape upfront. Call clear_shape when you are done.");
 	m.def("clear_shape", &xs3d::clear_shape, "Delete the data that was persisted by set_shape.");
 }
+
+#undef REGISTER_CALCULATE_AREA_FOR_TYPE
\ No newline at end of file
diff --git a/src/xs3d.hpp b/src/xs3d.hpp
index 0d28b45..94ec3ed 100644
--- a/src/xs3d.hpp
+++ b/src/xs3d.hpp
@@ -23,8 +23,9 @@ uint64_t _h(uint64_t a) {
 	return ((a+1) >> 1); 
 };
 
+template <typename LABEL>
 uint8_t compute_cube(
-	const uint8_t* binimg,
+	const LABEL* labels, const LABEL segid,
 	const uint64_t sx, const uint64_t sy, const uint64_t sz,
 	const uint64_t x, const uint64_t y, const uint64_t z
 ) {
@@ -36,14 +37,14 @@ uint8_t compute_cube(
     const uint64_t z_valid = (z < sz - 1);
 
 	return static_cast<uint8_t>(
-		(binimg[loc] > 0)
-		| ((x_valid && (binimg[loc+1] > 0)) << 1)
-		| ((y_valid && (binimg[loc+sx] > 0)) << 2)
-		| (((x_valid && y_valid) && (binimg[loc+sx+1] > 0)) << 3)
-		| ((z_valid && (binimg[loc+sxy] > 0)) << 4)
-		| (((x_valid && z_valid) && (binimg[loc+sxy+1] > 0)) << 5)
-		| (((y_valid && z_valid) && (binimg[loc+sxy+sx] > 0)) << 6)
-		| (((x_valid && y_valid && z_valid) && (binimg[loc+sxy+sx+1] > 0)) << 7)
+		(labels[loc] == segid)
+		| ((x_valid && (labels[loc+1] == segid)) << 1)
+		| ((y_valid && (labels[loc+sx] == segid)) << 2)
+		| (((x_valid && y_valid) && (labels[loc+sx+1] == segid)) << 3)
+		| ((z_valid && (labels[loc+sxy] == segid)) << 4)
+		| (((x_valid && z_valid) && (labels[loc+sxy+1] == segid)) << 5)
+		| (((y_valid && z_valid) && (labels[loc+sxy+sx] == segid)) << 6)
+		| (((x_valid && y_valid && z_valid) && (labels[loc+sxy+sx+1] == segid)) << 7)
 	);
 }
 
@@ -439,8 +440,9 @@ float calc_area_at_point_2x2x2(
 	return area;
 }
 
+template <typename LABEL>
 float calc_area_at_point(
-	const uint8_t* binimg,
+	const LABEL* labels, const LABEL segid,
 	std::vector<bool>& ccl,
 	const uint64_t sx, const uint64_t sy, const uint64_t sz,
 	const Vec3& cur, const Vec3& pos, 
@@ -504,7 +506,7 @@ float calc_area_at_point(
 				if (loc < 0 || loc >= voxels) {
 					continue;
 				}
-				else if (!binimg[loc]) {
+				else if (labels[loc] != segid) {
 					continue;
 				}
 
@@ -645,8 +647,9 @@ bool is_26_connected(
 	}
 }
 
+template <typename LABEL>
 float robust_calc_area_at_point_2x2x2(
-	const uint8_t* binimg,
+	const LABEL* labels, const LABEL segid,
 	std::vector<bool>& ccl,
 	const uint64_t sx, const uint64_t sy, const uint64_t sz,
 	const Vec3& cur, const Vec3& pos, 
@@ -684,7 +687,7 @@ float robust_calc_area_at_point_2x2x2(
 		ze = 0;		
 	}
 
-	const uint8_t center = compute_cube(binimg, sx, sy, sz, x, y, z);
+	const uint8_t center = compute_cube(labels, segid, sx, sy, sz, x, y, z);
 
 	for (int64_t zi = zs; zi <= ze; zi += 2) {
 		for (int64_t yi = ys; yi <= ye; yi += 2) {
@@ -701,13 +704,13 @@ float robust_calc_area_at_point_2x2x2(
 					(static_cast<uint64_t>(delta.y) >> 1) + _h(sy) * (static_cast<uint64_t>(delta.z) >> 1)
 				);
 
-				if (!binimg[loc] || ccl[ccl_loc]) {
+				if ((labels[loc] != segid) || ccl[ccl_loc]) {
 					continue;
 				}
 				
 				ccl[ccl_loc] = true;
 					
-				uint8_t candidate = compute_cube(binimg, sx, sy, sz, x + xi, y + yi, z + zi);
+				uint8_t candidate = compute_cube(labels, segid, sx, sy, sz, x + xi, y + yi, z + zi);
 
 				if (is_26_connected(center, candidate, xi, yi, zi)) {
 					subtotal += calc_area_at_point_2x2x2(
@@ -725,8 +728,9 @@ float robust_calc_area_at_point_2x2x2(
 	return subtotal;
 }
 
+template <typename LABEL>
 float robust_calc_area_at_point_2x2x2_persistent_data(
-	const uint8_t* binimg,
+	const LABEL* labels, const LABEL segid,
 	const uint64_t sx, const uint64_t sy, const uint64_t sz,
 	const Vec3& cur, const Vec3& pos, 
 	const Vec3& normal, const Vec3& anisotropy,
@@ -763,7 +767,7 @@ float robust_calc_area_at_point_2x2x2_persistent_data(
 		ze = 0;		
 	}
 
-	const uint8_t center = compute_cube(binimg, sx, sy, sz, x, y, z);
+	const uint8_t center = compute_cube(labels, segid, sx, sy, sz, x, y, z);
 	std::vector<uint8_t>& visited = persistent_data.visited;
 	const uint8_t color = persistent_data.color;
 
@@ -782,13 +786,13 @@ float robust_calc_area_at_point_2x2x2_persistent_data(
 					(static_cast<uint64_t>(delta.y) >> 1) + _h(sy) * (static_cast<uint64_t>(delta.z) >> 1)
 				);
 
-				if (!binimg[loc] || visited[visited_loc] == color) {
+				if ((labels[loc] != segid) || visited[visited_loc] == color) {
 					continue;
 				}
 				
 				visited[visited_loc] = color;
 				
-				const uint8_t candidate = compute_cube(binimg, sx, sy, sz, x + xi, y + yi, z + zi);
+				const uint8_t candidate = compute_cube(labels, segid, sx, sy, sz, x + xi, y + yi, z + zi);
 
 				if (is_26_connected(center, candidate, xi, yi, zi)) {
 					subtotal += calc_area_at_point_2x2x2(
@@ -806,8 +810,9 @@ float robust_calc_area_at_point_2x2x2_persistent_data(
 	return subtotal;
 }
 
+template <typename LABEL>
 std::tuple<float, uint8_t> cross_sectional_area_helper_2x2x2(
-	const uint8_t* binimg,
+	const LABEL* labels, const LABEL segid,
 	const uint64_t sx, const uint64_t sy, const uint64_t sz,
 	const Vec3& pos, // plane position
 	const Vec3& normal, // plane normal vector
@@ -902,7 +907,7 @@ std::tuple<float, uint8_t> cross_sectional_area_helper_2x2x2(
 			)
 		);
 
-		if (!binimg[loc]) {
+		if (labels[loc] != segid) {
 			continue;
 		}
 
@@ -949,8 +954,8 @@ std::tuple<float, uint8_t> cross_sectional_area_helper_2x2x2(
 			stack.push(downright);
 		}
 
-		total += robust_calc_area_at_point_2x2x2(
-			binimg, ccl,
+		total += robust_calc_area_at_point_2x2x2<LABEL>(
+			labels, segid, ccl,
 			sx, sy, sz,
 			cur, pos, normal, anisotropy,
 			pts, projections, inv_projections
@@ -960,8 +965,9 @@ std::tuple<float, uint8_t> cross_sectional_area_helper_2x2x2(
 	return std::make_tuple(total, contact);
 }
 
+template <typename LABEL>
 std::tuple<float, uint8_t> cross_sectional_area_helper_2x2x2_persistent_data(
-	const uint8_t* binimg,
+	const LABEL* labels, const LABEL segid,
 	const uint64_t sx, const uint64_t sy, const uint64_t sz,
 	const Vec3& pos, // plane position
 	const Vec3& normal, // plane normal vector
@@ -1056,7 +1062,7 @@ std::tuple<float, uint8_t> cross_sectional_area_helper_2x2x2_persistent_data(
 			)
 		);
 
-		if (!binimg[loc]) {
+		if (labels[loc] != segid) {
 			continue;
 		}
 
@@ -1104,7 +1110,7 @@ std::tuple<float, uint8_t> cross_sectional_area_helper_2x2x2_persistent_data(
 		}
 
 		total += robust_calc_area_at_point_2x2x2_persistent_data(
-			binimg,
+			labels, segid,
 			sx, sy, sz,
 			cur, pos, normal, anisotropy,
 			pts, projections, inv_projections
@@ -1114,9 +1120,9 @@ std::tuple<float, uint8_t> cross_sectional_area_helper_2x2x2_persistent_data(
 	return std::make_tuple(total, contact);
 }
 
-
+template <typename LABEL>
 float cross_sectional_area_helper(
-	const uint8_t* binimg,
+	const LABEL* labels, const LABEL segid,
 	const uint64_t sx, const uint64_t sy, const uint64_t sz,
 	const Vec3& pos, // plane position
 	const Vec3& normal, // plane normal vector
@@ -1208,7 +1214,7 @@ float cross_sectional_area_helper(
 			)
 		);
 
-		if (!binimg[loc]) {
+		if (labels[loc] != segid) {
 			continue;
 		}
 
@@ -1256,7 +1262,7 @@ float cross_sectional_area_helper(
 		}
 
 		total += calc_area_at_point(
-			binimg, ccl,
+			labels, segid, ccl,
 			sx, sy, sz,
 			cur, pos, normal, anisotropy,
 			pts, projections, inv_projections,
@@ -1302,8 +1308,9 @@ void clear_shape() {
 	persistent_data.clear();
 }
 
+template <typename LABEL>
 std::tuple<float, uint8_t> cross_sectional_area(
-	const uint8_t* binimg,
+	const LABEL* labels, const LABEL segid,
 	const uint64_t sx, const uint64_t sy, const uint64_t sz,
 	
 	const float px, const float py, const float pz,
@@ -1340,7 +1347,7 @@ std::tuple<float, uint8_t> cross_sectional_area(
 	if (loc < 0 || loc >= sx * sy * sz) {
 		return std::make_tuple(0.0, 0);
 	}
-	else if (!binimg[loc]) {
+	else if (labels[loc] != segid) {
 		return std::make_tuple(0.0, 0);
 	}
 
@@ -1349,15 +1356,15 @@ std::tuple<float, uint8_t> cross_sectional_area(
 	normal /= normal.norm();
 
 	if (use_persistent_data) {
-		return cross_sectional_area_helper_2x2x2_persistent_data(
-			binimg, 
+		return cross_sectional_area_helper_2x2x2_persistent_data<LABEL>(
+			labels, segid,
 			sx, sy, sz, 
 			pos, normal, anisotropy
 		);
 	}
 	else {
-		return cross_sectional_area_helper_2x2x2(
-			binimg, 
+		return cross_sectional_area_helper_2x2x2<LABEL>(
+			labels, segid,
 			sx, sy, sz, 
 			pos, normal, anisotropy
 		);
@@ -1416,7 +1423,7 @@ std::tuple<float*, uint8_t> cross_section(
 	normal /= normal.norm();
 
 	cross_sectional_area_helper(
-		binimg, 
+		binimg, /*segid=*/static_cast<uint8_t>(1),
 		sx, sy, sz, 
 		pos, normal, anisotropy,
 		contact, plane_visualization
diff --git a/xs3d/__init__.py b/xs3d/__init__.py
index e9775ad..b5da796 100644
--- a/xs3d/__init__.py
+++ b/xs3d/__init__.py
@@ -9,19 +9,21 @@
 
 
 def cross_sectional_area(
-  binimg:npt.NDArray[np.bool_],
+  labels:npt.NDArray[np.number],
   pos:POINT_T,
   normal:VECTOR_T,
   anisotropy:Optional[VECTOR_T] = None,
   return_contact:bool = False,
   slow_method:bool = False,
   use_persistent_data:bool = False,
+  segid:int|float = 1,
 ) -> Union[float, tuple[float, int]]:
   """
   Find the cross sectional area for a given binary image, 
   point, and normal vector.
 
-  binimg: a binary 2d or 3d numpy image (e.g. a bool datatype)
+  labels: a 2d or 3d numpy image
+  segid: which label to consider foreground 
   pos: the point in the image from which to extract the section
     must be an integer (it's an index into the image).
     e.g. [5,10,2]
@@ -58,15 +60,15 @@ def cross_sectional_area(
   """
   if anisotropy is None:
     anisotropy = (1.0, 1.0, 1.0)
-    if binimg.ndim == 2:
+    if labels.ndim == 2:
       anisotropy = (1.0, 1.0)
 
   pos = np.asarray(pos, dtype=np.float32)
   normal = np.asarray(normal, dtype=np.float32)
   anisotropy = np.asarray(anisotropy, dtype=np.float32)
 
-  if binimg.dtype != bool:
-    raise ValueError(f"A boolean image is required. Got: {binimg.dtype}")
+  if labels.dtype != bool:
+    raise ValueError(f"A boolean image is required. Got: {labels.dtype}")
 
   if np.any(anisotropy <= 0):
     raise ValueError(f"anisotropy values must be > 0. Got: {anisotropy}")
@@ -74,13 +76,13 @@ def cross_sectional_area(
   if np.all(normal == 0):
     raise ValueError("normal vector must not be a null vector (all zeros).")
 
-  binimg = np.asfortranarray(binimg)
+  labels = np.asfortranarray(labels)
 
-  if binimg.ndim == 2:
-    area, contact = cross_sectional_area_2d(binimg, pos, normal, anisotropy)
-  elif binimg.ndim == 3:
+  if labels.ndim == 2:
+    area, contact = cross_sectional_area_2d(labels, pos, normal, anisotropy)
+  elif labels.ndim == 3:
     area, contact = fastxs3d.area(
-      binimg.view(np.uint8),
+      labels, segid,
       pos, normal, anisotropy, 
       slow_method, use_persistent_data,
     )

From b394ae5ddd087b773608feab87248a3fe08eed45 Mon Sep 17 00:00:00 2001
From: William Silversmith <william.silversmith@gmail.com>
Date: Thu, 16 Oct 2025 20:09:07 -0400
Subject: [PATCH 2/4] fix: boolean images are no longer required

---
 xs3d/__init__.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/xs3d/__init__.py b/xs3d/__init__.py
index b5da796..f76d929 100644
--- a/xs3d/__init__.py
+++ b/xs3d/__init__.py
@@ -67,9 +67,6 @@ def cross_sectional_area(
   normal = np.asarray(normal, dtype=np.float32)
   anisotropy = np.asarray(anisotropy, dtype=np.float32)
 
-  if labels.dtype != bool:
-    raise ValueError(f"A boolean image is required. Got: {labels.dtype}")
-
   if np.any(anisotropy <= 0):
     raise ValueError(f"anisotropy values must be > 0. Got: {anisotropy}")
 

From 08badffb176c7e257c3584f25873e7858fbd5f01 Mon Sep 17 00:00:00 2001
From: William Silversmith <william.silversmith@gmail.com>
Date: Thu, 16 Oct 2025 20:12:08 -0400
Subject: [PATCH 3/4] fix: set setid to 1 only for binary images

---
 xs3d/__init__.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/xs3d/__init__.py b/xs3d/__init__.py
index f76d929..f9f27c3 100644
--- a/xs3d/__init__.py
+++ b/xs3d/__init__.py
@@ -16,7 +16,7 @@ def cross_sectional_area(
   return_contact:bool = False,
   slow_method:bool = False,
   use_persistent_data:bool = False,
-  segid:int|float = 1,
+  segid:Optional[Union[int,float]] = None,
 ) -> Union[float, tuple[float, int]]:
   """
   Find the cross sectional area for a given binary image, 
@@ -67,6 +67,12 @@ def cross_sectional_area(
   normal = np.asarray(normal, dtype=np.float32)
   anisotropy = np.asarray(anisotropy, dtype=np.float32)
 
+  if not np.issubdtype(labels.dtype, bool):
+    if segid is None:
+      raise ValueError(f"You must supply a segid for non-binary images.")
+  else:
+    segid = 1
+
   if np.any(anisotropy <= 0):
     raise ValueError(f"anisotropy values must be > 0. Got: {anisotropy}")
 

From 3bc493b3ba212bfd0e3de02e8134feabf57760af Mon Sep 17 00:00:00 2001
From: William Silversmith <william.silversmith@gmail.com>
Date: Thu, 16 Oct 2025 20:13:32 -0400
Subject: [PATCH 4/4] docs: update to note that 1 for binary images

---
 xs3d/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/xs3d/__init__.py b/xs3d/__init__.py
index f9f27c3..e457b37 100644
--- a/xs3d/__init__.py
+++ b/xs3d/__init__.py
@@ -23,7 +23,7 @@ def cross_sectional_area(
   point, and normal vector.
 
   labels: a 2d or 3d numpy image
-  segid: which label to consider foreground 
+  segid: which label to consider foreground (default 1 for boolean images)
   pos: the point in the image from which to extract the section
     must be an integer (it's an index into the image).
     e.g. [5,10,2]