From 653d53adde605c9225f67b4d78872f429635d051 Mon Sep 17 00:00:00 2001 From: William Silversmith Date: Sat, 22 Jun 2024 02:35:58 -0400 Subject: [PATCH] feat: support 2d arrays properly in cc3d.statistics --- automated_test.py | 45 ++++++++++++++++++- cc3d.pyx | 109 ++++++++++++++++++++++++++++++++++++++-------- 2 files changed, 136 insertions(+), 18 deletions(-) diff --git a/automated_test.py b/automated_test.py index 2d3977b..ff093ee 100644 --- a/automated_test.py +++ b/automated_test.py @@ -939,7 +939,7 @@ def test_voxel_graph_3d(): assert np.all(gt.T == graph) @pytest.mark.parametrize("order", ("C", "F")) -def test_statistics(order): +def test_statistics_3d(order): labels = np.zeros((123,128,125), dtype=np.uint8, order=order) labels[10:20,10:20,10:20] = 1 labels[40:50,40:50,40:51] = 2 @@ -978,6 +978,49 @@ def test_statistics(order): stats = cc3d.statistics(labels) assert np.all(stats["centroids"][0] == np.array([255.5,255.5,255.5])) +@pytest.mark.parametrize("order", ("C", "F")) +def test_statistics_2d(order): + labels = np.zeros((123,128), dtype=np.uint8, order=order) + labels[10:20,10:20] = 1 + labels[40:50,40:50] = 2 + + stats = cc3d.statistics(labels) + assert stats["voxel_counts"][1] == 100 + assert stats["voxel_counts"][2] == 10 * 10 + + assert np.all(np.isclose(stats["centroids"][1,:], [14.5,14.5])) + assert np.all(np.isclose(stats["centroids"][2,:], [44.5,44.5])) + + print(stats["bounding_boxes"]) + + assert np.all(stats["bounding_boxes"][0] == (slice(0,123), slice(0,128))) + assert np.all(stats["bounding_boxes"][1] == (slice(10,20), slice(10,20))) + assert np.all(stats["bounding_boxes"][2] == (slice(40,50), slice(40,50))) + + stats = cc3d.statistics(labels, no_slice_conversion=True) + print(stats["bounding_boxes"]) + assert np.all(stats["bounding_boxes"][0] == np.array([ 0, 122, 0, 127 ])) + assert np.all(stats["bounding_boxes"][1] == np.array([ 10, 19, 10, 19 ])) + assert np.all(stats["bounding_boxes"][2] == np.array([ 40, 49, 40, 49 ])) + + labels = np.zeros((1,1), dtype=np.uint8, order=order) + stats = cc3d.statistics(labels) + assert len(stats["voxel_counts"]) == 1 + assert stats["voxel_counts"][0] == 1 + + labels = np.zeros((0,1), dtype=np.uint8, order=order) + stats = cc3d.statistics(labels) + assert stats == { + "voxel_counts": None, + "bounding_boxes": None, + "centroids": None + } + + labels = np.zeros((512,512), dtype=np.uint8, order=order) + stats = cc3d.statistics(labels) + assert np.all(stats["centroids"][0] == np.array([255.5,255.5])) + + @pytest.mark.parametrize("order", ["C", "F"]) def test_statistics_big(order): labels = np.zeros((50,66000,1), dtype=np.uint8, order=order) diff --git a/cc3d.pyx b/cc3d.pyx index 4a9a735..41a8c35 100644 --- a/cc3d.pyx +++ b/cc3d.pyx @@ -665,26 +665,16 @@ def statistics( centroids: np.ndarray[float64] (N+1,3) } """ - while out_labels.ndim < 3: + while out_labels.ndim < 2: out_labels = out_labels[..., np.newaxis] if out_labels.dtype == bool: out_labels = out_labels.view(np.uint8) - return _statistics(out_labels, no_slice_conversion) - -@cython.cdivision(True) -@cython.boundscheck(False) -@cython.wraparound(False) -@cython.nonecheck(False) -def _statistics( - cnp.ndarray[UINT, ndim=3] out_labels, - native_bool no_slice_conversion -): cdef uint64_t voxels = out_labels.size; cdef uint64_t sx = out_labels.shape[0] cdef uint64_t sy = out_labels.shape[1] - cdef uint64_t sz = out_labels.shape[2] + cdef uint64_t sz = (out_labels.shape[2] if out_labels.ndim > 2 else 1) if voxels == 0: return { @@ -704,17 +694,102 @@ def _statistics( cdef cnp.ndarray[uint32_t] bounding_boxes32 if np.any(np.array([sx,sy,sz]) > np.iinfo(np.uint16).max): - bounding_boxes32 = np.zeros(6 * (N + 1), dtype=np.uint32) - return _statistics_helper(out_labels, no_slice_conversion, bounding_boxes32, N) + if out_labels.ndim == 2: + bounding_boxes32 = np.zeros(4 * (N + 1), dtype=np.uint32) + return _statistics_helper2d(out_labels, no_slice_conversion, bounding_boxes32, N) + else: + bounding_boxes32 = np.zeros(6 * (N + 1), dtype=np.uint32) + return _statistics_helper3d(out_labels, no_slice_conversion, bounding_boxes32, N) else: - bounding_boxes16 = np.zeros(6 * (N + 1), dtype=np.uint16) - return _statistics_helper(out_labels, no_slice_conversion, bounding_boxes16, N) + if out_labels.ndim == 2: + bounding_boxes32 = np.zeros(4 * (N + 1), dtype=np.uint32) + return _statistics_helper2d(out_labels, no_slice_conversion, bounding_boxes32, N) + else: + bounding_boxes16 = np.zeros(6 * (N + 1), dtype=np.uint16) + return _statistics_helper3d(out_labels, no_slice_conversion, bounding_boxes16, N) + +@cython.cdivision(True) +@cython.boundscheck(False) +@cython.wraparound(False) +@cython.nonecheck(False) +def _statistics_helper2d( + cnp.ndarray[UINT, ndim=2] out_labels, + native_bool no_slice_conversion, + cnp.ndarray[BBOX_T, ndim=1] bounding_boxes, + uint64_t N +): + cdef uint64_t voxels = out_labels.size; + cdef uint64_t sx = out_labels.shape[0] + cdef uint64_t sy = out_labels.shape[1] + + cdef cnp.ndarray[uint32_t] counts = np.zeros(N + 1, dtype=np.uint32) + cdef cnp.ndarray[double] centroids = np.zeros(2 * (N + 1), dtype=np.float64) + + cdef BBOX_T x = 0 + cdef BBOX_T y = 0 + + cdef uint64_t label = 0 + + bounding_boxes[::2] = np.iinfo(bounding_boxes.dtype).max + + if out_labels.flags.f_contiguous: + for y in range(sy): + for x in range(sx): + label = out_labels[x,y] + counts[label] += 1 + bounding_boxes[4 * label + 0] = min(bounding_boxes[4 * label + 0], x) + bounding_boxes[4 * label + 1] = max(bounding_boxes[4 * label + 1], x) + bounding_boxes[4 * label + 2] = min(bounding_boxes[4 * label + 2], y) + bounding_boxes[4 * label + 3] = max(bounding_boxes[4 * label + 3], y) + centroids[2 * label + 0] += x + centroids[2 * label + 1] += y + else: + for x in range(sx): + for y in range(sy): + label = out_labels[x,y] + counts[label] += 1 + bounding_boxes[4 * label + 0] = min(bounding_boxes[4 * label + 0], x) + bounding_boxes[4 * label + 1] = max(bounding_boxes[4 * label + 1], x) + bounding_boxes[4 * label + 2] = min(bounding_boxes[4 * label + 2], y) + bounding_boxes[4 * label + 3] = max(bounding_boxes[4 * label + 3], y) + centroids[2 * label + 0] += x + centroids[2 * label + 1] += y + + for label in range(N+1): + if counts[label] == 0: + centroids[2 * label + 0] = float('NaN') + centroids[2 * label + 1] = float('NaN') + else: + centroids[2 * label + 0] /= counts[label] + centroids[2 * label + 1] /= counts[label] + + bbxes = bounding_boxes.reshape((N+1,4)) + + output = { + "voxel_counts": counts, + "bounding_boxes": bbxes, + "centroids": centroids.reshape((N+1,2)), + } + + if no_slice_conversion: + return output + + slices = [] + for xs, xe, ys, ye in bbxes: + if xs < voxels and ys < voxels: + slices.append((slice(xs, int(xe+1)), slice(ys, int(ye+1)))) + else: + slices.append(None) + + output["bounding_boxes"] = slices + + return output @cython.cdivision(True) @cython.boundscheck(False) @cython.wraparound(False) @cython.nonecheck(False) -def _statistics_helper( +def _statistics_helper3d( cnp.ndarray[UINT, ndim=3] out_labels, native_bool no_slice_conversion, cnp.ndarray[BBOX_T, ndim=1] bounding_boxes,