From 12d683258b6f78391afced3102d6826c8579bcbd Mon Sep 17 00:00:00 2001 From: Naman Gera Date: Mon, 10 Jul 2023 13:51:26 +0100 Subject: [PATCH 1/2] Make distortion correction method work only for 3d data --- httomolibgpu/prep/alignment.py | 6 +-- tests/test_prep/test_alignment.py | 73 ++++++++++++++++++++++--------- 2 files changed, 55 insertions(+), 24 deletions(-) diff --git a/httomolibgpu/prep/alignment.py b/httomolibgpu/prep/alignment.py index 484fe9a3..f1ef756b 100644 --- a/httomolibgpu/prep/alignment.py +++ b/httomolibgpu/prep/alignment.py @@ -254,9 +254,9 @@ def distortion_correction_proj_discorpy( cp.ndarray 3D array. Distortion-corrected image(s). """ - # Check if it's a stack of 2D images, or only a single 2D image - if len(data.shape) == 2: - data = cp.expand_dims(data, axis=0) + # Check if it's a stack of 2D images only + if data.ndim != 3: + raise ValueError("Input data must be a stack of 2D images only") # Get info from metadata txt file xcenter, ycenter, list_fact = _load_metadata_txt(metadata_path) diff --git a/tests/test_prep/test_alignment.py b/tests/test_prep/test_alignment.py index 1499c771..20385ce7 100644 --- a/tests/test_prep/test_alignment.py +++ b/tests/test_prep/test_alignment.py @@ -16,11 +16,11 @@ @cp.testing.gpu @pytest.mark.parametrize( - "image, max_value, mean_value", + "image, mean_value, mean_sum, sum", [ - ("dot_pattern_03.tif", 255, 200.16733869461675), - ("peppers.tif", 228, 95.51871109008789), - ("cameraman.tif", 254, 122.2400016784668), + ("dot_pattern_03.tif", 200.16733869461675, 20016.733869461674, 214366729300), + ("peppers.tif", 95.51871109008789, 9551.871109008789, 2503965700), + ("cameraman.tif", 122.2400016784668, 12224.00016784668, 3204448300), ], ids=["dot_pattern_03", "peppers", "cameraman"], ) @@ -34,24 +34,42 @@ def test_correct_distortion( ensure_clean_memory, image, implementation, - max_value, mean_value, + mean_sum, + sum, ): distortion_coeffs_path = os.path.join( distortion_correction_path, "distortion-coeffs.txt" ) path = os.path.join(distortion_correction_path, image) - im_host = imread(path) - im = cp.asarray(im_host) + im_host = np.asarray(imread(path)) + im_host = np.expand_dims(im_host, axis=0) + # replicate into a stack of images + im_stack = cp.asarray(cp.tile(im_host, (100, 1, 1))) + + preview = { + "starts": [0, 0], + "stops": [im_stack.shape[0], im_stack.shape[1]], + "steps": [1, 1], + } + corrected_data = implementation(im_stack, distortion_coeffs_path, preview).get() - preview = {"starts": [0, 0], "stops": [im.shape[0], im.shape[1]], "steps": [1, 1]} - corrected_data = implementation(im, distortion_coeffs_path, preview).get() - assert_allclose(np.mean(corrected_data), mean_value) - assert np.max(corrected_data) == max_value + assert_allclose(np.mean(corrected_data, axis=(1, 2)).sum(), mean_sum, rtol=1e-6) + assert np.sum(corrected_data) == sum assert corrected_data.dtype == np.uint8 + assert corrected_data.ndim == 3 + + +def test_correct_distortion_1d_raises(distortion_correction_path, ensure_clean_memory): + with pytest.raises(ValueError): + preview = {"starts": [0, 0], "stops": [1, 1], "steps": [1, 1]} + distortion_coeffs_path = os.path.join( + distortion_correction_path, "distortion-coeffs.txt" + ) + distortion_correction_proj_discorpy(cp.ones(1), distortion_coeffs_path, preview) @pytest.mark.parametrize("stack_size", [20, 50, 100]) @@ -60,7 +78,9 @@ def test_correct_distortion( [distortion_correction_proj, distortion_correction_proj_discorpy], ids=["cupy", "tomopy"], ) -def test_distortion_correction_meta(distortion_correction_path, stack_size, implementation): +def test_distortion_correction_meta( + distortion_correction_path, stack_size, implementation +): distortion_coeffs_path = os.path.join( distortion_correction_path, "distortion-coeffs.txt" ) @@ -70,20 +90,31 @@ def test_distortion_correction_meta(distortion_correction_path, stack_size, impl im_host = np.expand_dims(im_host, axis=0) # replicate into a stack of images im_stack = cp.asarray(np.tile(im_host, (stack_size, 1, 1))) + hook = MaxMemoryHook(im_stack.size * im_stack.itemsize) - preview = {"starts": [0, 0], "stops": [im_stack.shape[1], im_stack.shape[2]], "steps": [1, 1]} - + preview = { + "starts": [0, 0], + "stops": [im_stack.shape[1], im_stack.shape[2]], + "steps": [1, 1], + } + with hook: implementation(im_stack, distortion_coeffs_path, preview) - + # make sure estimator function is within range (80% min, 100% max) max_mem = hook.max_mem actual_slices = im_stack.shape[0] estimated_slices, dtype_out, output_dims = implementation.meta.calc_max_slices( - 0, (im_stack.shape[1], im_stack.shape[2]), - im_stack.dtype, max_mem, - metadata_path=distortion_coeffs_path, preview=preview) + 0, + (im_stack.shape[1], im_stack.shape[2]), + im_stack.dtype, + max_mem, + metadata_path=distortion_coeffs_path, + preview=preview, + ) assert estimated_slices <= actual_slices - assert estimated_slices / actual_slices >= 0.8 - - assert implementation.__name__ in method_registry['httomolibgpu']['prep']['alignment'] + assert estimated_slices / actual_slices >= 0.8 + + assert ( + implementation.__name__ in method_registry["httomolibgpu"]["prep"]["alignment"] + ) From 2ba01d2d1e562293ff365cd276c8e76fd3f9280e Mon Sep 17 00:00:00 2001 From: Naman Gera Date: Mon, 10 Jul 2023 15:03:54 +0100 Subject: [PATCH 2/2] Add performance test --- tests/test_prep/test_alignment.py | 47 ++++++++++++++++++++++++++++--- 1 file changed, 43 insertions(+), 4 deletions(-) diff --git a/tests/test_prep/test_alignment.py b/tests/test_prep/test_alignment.py index 20385ce7..423e259d 100644 --- a/tests/test_prep/test_alignment.py +++ b/tests/test_prep/test_alignment.py @@ -1,16 +1,18 @@ import os +import time import cupy as cp import numpy as np import pytest +from cupy.cuda import nvtx +from imageio.v2 import imread +from numpy.testing import assert_allclose + +from httomolibgpu import method_registry from httomolibgpu.prep.alignment import ( distortion_correction_proj, distortion_correction_proj_discorpy, ) -from httomolibgpu import method_registry -from imageio.v2 import imread -from numpy.testing import assert_allclose - from tests import MaxMemoryHook @@ -118,3 +120,40 @@ def test_distortion_correction_meta( assert ( implementation.__name__ in method_registry["httomolibgpu"]["prep"]["alignment"] ) + + +@cp.testing.gpu +@pytest.mark.perf +def test_distortion_correction_performance( + distortion_correction_path, ensure_clean_memory +): + dev = cp.cuda.Device() + distortion_coeffs_path = os.path.join( + distortion_correction_path, "distortion-coeffs.txt" + ) + + image_path = os.path.join(distortion_correction_path, "dot_pattern_03.tif") + im_host = np.asarray(imread(image_path)) + im_host = np.expand_dims(im_host, axis=0) + # replicate into a stack of images + im_stack = cp.asarray(np.tile(im_host, (150, 3, 2))) + + preview = { + "starts": [0, 0], + "stops": [im_stack.shape[0], im_stack.shape[1]], + "steps": [1, 1], + } + + # cold run first + distortion_correction_proj_discorpy(im_stack, distortion_coeffs_path, preview) + dev.synchronize() + + start = time.perf_counter_ns() + nvtx.RangePush("Core") + for _ in range(10): + distortion_correction_proj_discorpy(im_stack, distortion_coeffs_path, preview) + nvtx.RangePop() + dev.synchronize() + duration_ms = float(time.perf_counter_ns() - start) * 1e-6 / 10 + + assert "performance in ms" == duration_ms