Merge pull request #73 from Hendrik-code/speed_extract_cc_overhaul

Speed extract cc overhaul

Author: Hendrik-code

TPTBox/core/nii_wrapper.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import random
+import timeit
 from collections.abc import Callable
 from copy import deepcopy
 from time import perf_counter
@@ -8,14 +10,35 @@
 from tqdm import tqdm
 
 
-def speed_test_input(inp, functions: list[Callable], assert_equal_function: Callable | None = None, *args, **kwargs):
+def speed_test_input(
+    inp,
+    functions: list[Callable],
+    assert_equal_function: Callable | None = None,
+    print_output: bool = False,
+    *args,
+    **kwargs,
+):
     time_measures = {}
     outs = {}
+    random.shuffle(functions)
     for f in functions:
-        start = perf_counter()
         input_copy = deepcopy(inp)
         out = f(*input_copy, *args, **kwargs) if isinstance(input_copy, (tuple, list)) else f(input_copy, *args, **kwargs)
-        time = perf_counter() - start
+        if print_output:
+            print(f.__name__, out)
+
+        # def f_test():
+        #    f(*input_copy, *args, **kwargs) if isinstance(input_copy, (tuple, list)) else f(input_copy, *args, **kwargs)
+
+        time = timeit.timeit(
+            "f(*input_copy, *args, **kwargs) if isinstance(input_copy, (tuple, list)) else f(input_copy, *args, **kwargs)",
+            # "f_test()",
+            # setup="def f_test(): f(*input_copy, *args, **kwargs) if isinstance(input_copy, (tuple, list)) else f(input_copy, *args, **kwargs)",
+            number=5,
+            globals=locals(),
+        )
+
+        # time = perf_counter() - start
         outs[f.__name__] = out
         time_measures[f.__name__] = time
 
@@ -38,22 +61,27 @@ def speed_test(
     # print first iteration
     print()
     print("Print first speed test")
-    start = perf_counter()
-    first_input = get_input_func()
-    time = perf_counter() - start
-    time_sums["input_function"].append(time)
-    for f in functions:
-        input_copy = deepcopy(first_input)
-        out = f(*input_copy, *args, **kwargs) if isinstance(input_copy, (tuple, list)) else f(input_copy, *args, **kwargs)
-        print(f.__name__, out)
 
-    for _ in tqdm(range(repeats)):
+    for repeat_idx in tqdm(range(repeats)):
+        start = perf_counter()
         inp = get_input_func()
-        time_measures = speed_test_input(inp, *args, functions=functions, assert_equal_function=assert_equal_function, **kwargs)
+        time = perf_counter() - start
+        time_sums["input_function"].append(time)
+        time_measures = speed_test_input(
+            inp,
+            *args,
+            functions=functions,
+            assert_equal_function=assert_equal_function,
+            print_output=repeat_idx == 0,
+            **kwargs,
+        )
         for k, v in time_measures.items():
             if k not in time_sums:
                 time_sums[k] = []
             time_sums[k].append(v)
 
-    for k, v in time_sums.items():
-        print(k, "\t", round(sum(v) / repeats, ndigits=6), "+-", round(np.std(v), ndigits=6))
+    times_sorted = dict(sorted(time_sums.items(), key=lambda x: sum(x[1]) / repeats))
+    for idx, (k, v) in enumerate(times_sorted.items()):
+        print(idx + 1, ".\t", round(sum(v) / repeats, ndigits=6), "+-", round(np.std(v), ndigits=6), "\t", k)
+    # for k, v in time_sums.items():
+    #    print(k, "\t", round(sum(v) / repeats, ndigits=6), "+-", round(np.std(v), ndigits=6))

TPTBox/core/np_utils.py

@@ -10,7 +10,6 @@
 
     from TPTBox.core.np_utils import (
         _to_labels,
-        np_approx_center_of_mass,
         np_bbox_binary,
         np_bounding_boxes,
         np_center_of_mass,
@@ -26,8 +25,8 @@ def get_nii_array():
         num_points = random.randint(1, 30)
         nii, points, orientation, sizes = get_nii(x=(140, 140, 150), num_point=num_points)
         # nii.map_labels_({1: -1}, verbose=False)
-        arr = nii.get_seg_array().astype(int)
-        arr[arr == 1] = -1
+        arr = nii.get_seg_array().astype(np.uint8)
+        # arr[arr == 1] = -1
         arr_r = arr.copy()
         return arr_r
 
@@ -39,64 +38,30 @@ def center_of_mass_one(arr: np.ndarray):
         return coms
 
     def center_of_mass_(arr: np.ndarray):
-        cc_label_set = np.unique(arr)
+        cc_label_set = np_unique(arr)
         coms = {}
         for c in cc_label_set:
             if c == 0:
                 continue
-            c_l = arr.copy()
-            c_l[c_l != c] = 0
-            com = center_of_mass(c_l)
+            com = center_of_mass(arr == c)
             coms[c] = com
         return coms
 
     def bbox_(arr: np.ndarray):
-        cc_label_set = np.unique(arr)
+        cc_label_set = np_unique(arr)
         coms = {}
         for c in cc_label_set:
             if c == 0:
                 continue
-            c_l = arr.copy()
-            c_l[c_l != c] = 0
-            com = np_bbox_binary(c_l)
+            com = np_bbox_binary(arr == c)
             coms[c] = com
         return coms
 
-    # def cc3d_volume(arr: np.ndarray):
-    #    volumes = dict(enumerate(cc3dstats(arr)["voxel_counts"]))
-    #    volumes.pop(0)
-    #    return volumes
-
-    # def cc3d_countnonzero(arr: np.ndarray):
-    #    return sum(cc3d_volume(arr).values())
-
-    # def cc3d_bbox(arr: np.ndarray):
-    #    return dict(enumerate(cc3dstats(arr)["bounding_boxes"]))
-
-    # def cc3d_unique(arr: np.ndarray) -> list[int]:
-    #    return [i for i, v in cc3d_volume(arr).items() if v > 0]
-
-    # a = np.ones((100, 100, 50), dtype=np.uint16)
-    # a[0, 0, 0] = 0
-    # print(type(a[0, 0, 0]))
-    # print(cc3d_volume(a))
-
-    arr = get_nii_array()
-    # print(np_unique_withoutzero(arr))
-    # print("npunique", np.unique(arr))
-    # print()
-    # print(center_of_mass_(arr))
-    # print()
-    # print(np_center_of_mass(arr))
-    print()
-    print(np_unique(arr))
-    print(np.unique(arr))
-
     speed_test(
         repeats=50,
         get_input_func=get_nii_array,
-        functions=[np_unique, np.unique],
-        assert_equal_function=lambda x, y: True,  # np.all([x[i] == y[i] for i in range(len(x))]),  # noqa: ARG005
+        functions=[cc3d_com, center_of_mass_],
+        assert_equal_function=lambda x, y: np.all([x[i][0] == y[i][0] for i in x.keys()]),  # noqa: ARG005
         # np.all([x[i] == y[i] for i in range(len(x))])
     )
     # print(time_measures)

TPTBox/tests/speedtests/speedtest.py

@@ -25,12 +25,12 @@ def get_nii_array():
         return arr
 
     def normal(arr):
-        return cc3dstatistics(arr)
+        return cc3dstatistics(arr, use_crop=False)
 
     def crop(arr):
         crop = np_bbox_binary(arr)
         arr = arr[crop]
-        return cc3dstatistics(arr)
+        return cc3dstatistics(arr, use_crop=False)
 
     speed_test(
         repeats=50,

TPTBox/tests/speedtests/speedtest_cc3d.py

@@ -22,7 +22,7 @@
 
     def get_nii_array():
         num_points = random.randint(50, 51)
-        nii, points, orientation, sizes = get_nii(x=(300, 300, 300), num_point=num_points)
+        nii, points, orientation, sizes = get_nii(x=(100, 100, 100), num_point=num_points)
         # nii.map_labels_({1: -1}, verbose=False)
         arr = nii.get_seg_array().astype(np.uint8)
         # arr[arr == 1] = -1
@@ -32,6 +32,9 @@ def get_nii_array():
     def np_naive_cc(arr: np.ndarray):
         return np_connected_components(arr)[0][1]
 
+    def np_naive_cc_extract(arr: np.ndarray):
+        return np_connected_components(arr, use_extract2=True)[0][1]
+
     def np_naive_cc_gcrop(arr: np.ndarray):
         crop = np_bbox_binary(arr)
         arrc = arr[crop]
@@ -72,31 +75,33 @@ def np_crop_cc(arr: np.ndarray):
             subreg_cc_n[subreg] = n
         return subreg_cc[1]  # , subreg_cc_n
 
-    def np_cc_once(arr: np.ndarray):
-        # call cc once, then relabel
-        connectivity = 3
-        connectivity = min((connectivity + 1) * 2, 8) if arr.ndim == 2 else 6 if connectivity == 1 else 18 if connectivity == 2 else 26
-
-        labels: list[int] = np_unique(arr)
-
-        subreg_cc = {}
-        subreg_cc_n = {}
-        crop = np_bbox_binary(arr)
-        arrc = arr[crop]
-        zarr = np.zeros((len(labels), *arr.shape), dtype=arr.dtype)
-
-        labels_out = connected_components(arrc, connectivity=connectivity, return_N=False)
-        for sidx, subreg in enumerate(labels):  # type:ignore
-            img_subreg = np_extract_label(arrc, subreg, inplace=False)
-            lcrop = np_bbox_binary(img_subreg)
-            img_subregc = img_subreg[lcrop]
-            img_subreg[lcrop] = labels_out[lcrop] * img_subregc
-
-            arrcc = zarr[sidx]
-            arrcc[crop] = img_subreg
-            subreg_cc[subreg] = arrcc
-            subreg_cc_n[subreg] = len(np_unique_withoutzero(img_subreg[lcrop]))
-        return subreg_cc[1]  # , subreg_cc_n
+    # def np_cc_once(arr: np.ndarray):
+    #    # call cc once, then relabel
+    #    connectivity = 3
+    #    connectivity = min((connectivity + 1) * 2, 8) if arr.ndim == 2 else 6 if connectivity == 1 else 18 if connectivity == 2 else 26
+    #
+    #    labels: list[int] = np_unique(arr)
+    #
+    #    subreg_cc = {}
+    #    subreg_cc_n = {}
+    #    crop = np_bbox_binary(arr)
+    #    arrc = arr[crop]
+    #    zarr = np.zeros((len(labels), *arr.shape), dtype=arr.dtype)
+    #
+    #    labels_out = connected_components(arrc, connectivity=connectivity, return_N=False)
+    #    for sidx, subreg in enumerate(labels):  # type:ignore
+    #        arrcc[crop][np.logical_and()]
+    #        # arr[s == subreg]
+    #        # img_subreg = np_extract_label(arrc, subreg, inplace=False)
+    #        # lcrop = np_bbox_binary(img_subreg)
+    #        img_subregc = img_subreg[lcrop]
+    #        img_subreg[lcrop] = labels_out[lcrop] * img_subregc
+    #
+    #        arrcc = zarr[sidx]
+    #        arrcc[crop] = img_subreg
+    #        subreg_cc[subreg] = arrcc
+    #        subreg_cc_n[subreg] = len(np_unique_withoutzero(img_subreg[lcrop]))
+    #    return subreg_cc[1]  # , subreg_cc_n
 
     def np_cc_once_lcrop(arr: np.ndarray):
         # call cc once, then relabel
@@ -125,9 +130,9 @@ def np_cc_once_lcrop(arr: np.ndarray):
         return subreg_cc[1]  # , subreg_cc_n
 
     speed_test(
-        repeats=10,
+        repeats=50,
         get_input_func=get_nii_array,
-        functions=[np_naive_cc, np_naive_cc_gcrop, np_cc_once],
+        functions=[np_naive_cc, np_naive_cc_extract],
         assert_equal_function=lambda x, y: np.count_nonzero(x) == np.count_nonzero(y),
         # np.all([x[i] == y[i] for i in range(x.shape[0])]),  # noqa: ARG005
         # np.all([x[i] == y[i] for i in range(len(x))])

TPTBox/tests/speedtests/speedtest_cc3d_crop.py

@@ -0,0 +1,82 @@
+if __name__ == "__main__":
+    # speed test dilation
+    import random
+
+    import numpy as np
+    from cc3d import statistics as cc3dstats
+    from scipy.ndimage import center_of_mass
+
+    from TPTBox.core.nii_wrapper import NII
+    from TPTBox.core.np_utils import (
+        _connected_components,
+        _to_labels,
+        np_bbox_binary,
+        np_calc_overlapping_labels,
+        np_connected_components,
+        np_connected_components_per_label,
+        np_connected_components_per_label2,
+        np_extract_label,
+        np_unique,
+        np_unique_withoutzero,
+    )
+    from TPTBox.tests.speedtests.speedtest import speed_test
+    from TPTBox.tests.test_utils import get_nii
+
+    def get_nii_array():
+        num_points = random.randint(10, 31)
+        nii, points, orientation, sizes = get_nii(x=(300, 300, 300), num_point=num_points)
+        # nii.map_labels_({1: -1}, verbose=False)
+        arr = nii.get_seg_array().astype(np.uint8)
+        # arr[arr == 1] = -1
+        # arr_r = arr.copy()
+        return arr
+
+    def np_cc_labelwise1(arr: np.ndarray):
+        return np_connected_components_per_label(arr)[0][1]
+
+    def np_cc_labelwise2(arr: np.ndarray):
+        return np_connected_components_per_label2(arr)[0][1]
+
+    def np_cc_labelwise2crop(arr: np.ndarray):
+        return np_connected_components_per_label2(arr, use_crop=True)[0][1]
+
+    # def np_cc_once(arr: np.ndarray):
+    #    # call cc once, then relabel
+    #    connectivity = 3
+    #    connectivity = min((connectivity + 1) * 2, 8) if arr.ndim == 2 else 6 if connectivity == 1 else 18 if connectivity == 2 else 26
+    #
+    #    labels: list[int] = np_unique(arr)
+    #
+    #    subreg_cc = {}
+    #    subreg_cc_n = {}
+    #    crop = np_bbox_binary(arr)
+    #    arrc = arr[crop]
+    #    zarr = np.zeros((len(labels), *arr.shape), dtype=arr.dtype)
+    #
+    #    labels_out = connected_components(arrc, connectivity=connectivity, return_N=False)
+    #    for sidx, subreg in enumerate(labels):  # type:ignore
+    #        arrcc[crop][np.logical_and()]
+    #        # arr[s == subreg]
+    #        # img_subreg = np_extract_label(arrc, subreg, inplace=False)
+    #        # lcrop = np_bbox_binary(img_subreg)
+    #        img_subregc = img_subreg[lcrop]
+    #        img_subreg[lcrop] = labels_out[lcrop] * img_subregc
+    #
+    #        arrcc = zarr[sidx]
+    #        arrcc[crop] = img_subreg
+    #        subreg_cc[subreg] = arrcc
+    #        subreg_cc_n[subreg] = len(np_unique_withoutzero(img_subreg[lcrop]))
+    #    return subreg_cc[1]  # , subreg_cc_n
+
+    speed_test(
+        repeats=50,
+        get_input_func=get_nii_array,
+        functions=[
+            np_cc_labelwise1,
+            np_cc_labelwise2,
+        ],
+        assert_equal_function=lambda x, y: True,  # np.array_equal(x, y),  # noqa: ARG005
+        # np.all([x[i] == y[i] for i in range(x.shape[0])]),  # noqa: ARG005
+        # np.all([x[i] == y[i] for i in range(len(x))])
+    )
+    # print(time_measures)