Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Multiscale filtering functionality #16

Merged
merged 28 commits into from
Jul 3, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
28 commits
Select commit Hold shift + click to select a range
ba97fd2
Added binning/debinning functions from bm3d package.
dmitry-ganyushin Jun 18, 2024
69b1926
Added test scripts and test data.
Jun 18, 2024
91604b5
Excluded failed due to changed API tests
Jun 20, 2024
4ed770a
Temporary disabled tests
Jun 21, 2024
aaf4e6f
Merge branch 'refs/heads/next' into bin
Jun 21, 2024
f8b1524
Merged with next
Jun 21, 2024
c5f4b06
Adapted example
Jun 21, 2024
6fca7ce
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 27, 2024
6c32366
Fixed issues in the review.
Jun 28, 2024
98094e9
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 1, 2024
47ec0b9
Addressed revision comments
Jul 2, 2024
4b78bd0
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 2, 2024
230272a
Added function parameters.
Jul 2, 2024
ff053b3
Merge remote-tracking branch 'origin/bin' into bin
Jul 2, 2024
1823b25
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 2, 2024
3f6f65a
Disabled test
dmitry-ganyushin Jul 2, 2024
5e7c593
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 2, 2024
94ae649
Disabled test
dmitry-ganyushin Jul 2, 2024
f5d32af
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 2, 2024
8bf61ec
Adjusted test
Jul 3, 2024
5519360
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 3, 2024
98f1153
Adjusted test
Jul 3, 2024
65d7b7d
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 3, 2024
8449ba5
Adjusted test
Jul 3, 2024
0e9fa09
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 3, 2024
9d38233
Added cuda decorator
Jul 3, 2024
6f7b2fa
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 3, 2024
972fa85
Converted test to python notebook example.
dmitry-ganyushin Jul 3, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 60 additions & 0 deletions notebooks/example_ms.ipynb
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"from bm3dornl.bm3d import bm3d_ring_artifact_removal_ms\n",
"from bm3dornl.denoiser_gpu import memory_cleanup\n",
"\n",
"with open(\"../../tests/bm3dornl-data/sino.npy\", \"rb\") as f:\n",
" sino_noisy = np.load(f)\n",
"\n",
"memory_cleanup()\n",
"\n",
"block_matching_kwargs: dict = {\n",
" \"patch_size\": (8, 8),\n",
" \"stride\": 3,\n",
" \"background_threshold\": 0.0,\n",
" \"cut_off_distance\": (64, 64),\n",
" \"num_patches_per_group\": 32,\n",
" \"padding_mode\": \"circular\",\n",
"}\n",
"filter_kwargs: dict = {\n",
" \"filter_function\": \"fft\",\n",
" \"shrinkage_factor\": 3e-2,\n",
"}\n",
"kwargs = {\n",
" \"mode\": \"simple\",\n",
" \"k\": 4,\n",
" \"block_matching_kwargs\": block_matching_kwargs,\n",
" \"filter_kwargs\": filter_kwargs,\n",
"}\n",
"\n",
"sino_bm3dornl = bm3d_ring_artifact_removal_ms(\n",
" sinogram=sino_noisy,\n",
" **kwargs,\n",
")\n",
"\n",
"\n",
"fig, axs = plt.subplots(1, 2, figsize=(12, 4))\n",
"axs[0].imshow(sino_noisy, cmap=\"gray\")\n",
"axs[0].set_title(\"Noisy sinogram\")\n",
"axs[1].imshow(sino_bm3dornl, cmap=\"gray\")\n",
"axs[1].set_title(\"BM3D denoised sinogram\")\n",
"plt.show()\n"
]
}
],
"metadata": {
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
56 changes: 38 additions & 18 deletions src/bm3dornl/bm3d.py
Original file line number Diff line number Diff line change
Expand Up @@ -739,26 +739,46 @@ def bm3d_ring_artifact_removal_ms(
filter_kwargs=filter_kwargs,
)

# step 1: create a list of binned sinograms
binned_sinos = horizontal_binning(sinogram, k=k)
# reverse the list
binned_sinos = binned_sinos[::-1]

# step 2: estimate the noise level from the coarsest sinogram, then working back to the original sinogram
noise_estimate = None
for i in range(len(binned_sinos)):
logging.info(f"Processing binned sinogram {i+1} of {len(binned_sinos)}")
sino = binned_sinos[i]
sino_star = (
sino if i == 0 else sino - horizontal_debinning(noise_estimate, sino)
denoised_sino = None
# Make a copy of an original sinogram
sino_orig = horizontal_binning(sino_star, 1, dim=0)
binned_sinos_orig = [np.copy(sino_orig)]

# Contains upscaled denoised sinograms
binned_sinos = [np.zeros(0)]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why creating an np.zeros(0) here? Are we trying to do np.empty or something else?


# Bin horizontally
for i in range(0, k):
binned_sinos_orig.append(
horizontal_binning(binned_sinos_orig[-1], fac=2, dim=1)
)
binned_sinos.append(np.zeros(0))

if i < len(binned_sinos) - 1:
noise_estimate = sino - bm3d_ring_artifact_removal(
sino_star,
mode=mode,
block_matching_kwargs=block_matching_kwargs,
filter_kwargs=filter_kwargs,
binned_sinos[-1] = binned_sinos_orig[-1]

for i in range(k, -1, -1):
logging.info(f"Processing binned sinogram {i + 1} of {k}")
# Denoise binned sinogram
denoised_sino = bm3d_ring_artifact_removal(
binned_sinos[i],
mode=mode,
block_matching_kwargs=block_matching_kwargs,
filter_kwargs=filter_kwargs,
)
# For iterations except the last, create the next noisy image with a finer scale residual
if i > 0:
debinned_sino = horizontal_debinning(
denoised_sino - binned_sinos_orig[i],
binned_sinos_orig[i - 1].shape[1],
2,
30,
Comment on lines +773 to +774
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hard-coded magic numbers should be avoided, or at least provide some comments on the reasoning.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added function parameters

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you explain a bit on the decision on selecting 2 and 30 here? This is to help the future developer so that they don't have to guess why these number are chosen.

dim=1,
)
binned_sinos[i - 1] = binned_sinos_orig[i - 1] + debinned_sino

# residual
sino_star = sino_star + horizontal_debinning(
denoised_sino - sino_orig, sino_star.shape[0], fac=1, n_iter=30, dim=0
)

return sino_star
161 changes: 104 additions & 57 deletions src/bm3dornl/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@
"""Utility functions for BM3DORNL."""

import numpy as np
from scipy.interpolate import RectBivariateSpline
from numba import njit
from scipy.signal import convolve2d
from scipy.interpolate import interp1d


@njit
dmitry-ganyushin marked this conversation as resolved.
Show resolved Hide resolved
Expand Down Expand Up @@ -43,84 +44,130 @@ def is_within_threshold(
return np.linalg.norm(ref_patch - cmp_patch) <= intensity_diff_threshold


def horizontal_binning(Z: np.ndarray, k: int = 0) -> list[np.ndarray]:
def create_array(base_arr: np.ndarray, h: int, dim: int):
"""
Horizontal binning of the image Z into a list of k images.
Create a padded and convolved array used in both binning and debinning.
Parameters
----------
base_arr: Input array
h: bin count
dim: bin dimension (0 or 1)
Returns
-------
resulting array
"""
mod_pad = h - ((base_arr.shape[dim] - 1) % h) - 1
if dim == 0:
pads = ((0, mod_pad), (0, 0))
kernel = np.ones((h, 1))
else:
pads = ((0, 0), (0, mod_pad))
kernel = np.ones((1, h))

return convolve2d(np.pad(base_arr, pads, "symmetric"), kernel, "same", "fill")


def horizontal_binning(Z: np.ndarray, fac: int = 2, dim: int = 1) -> np.ndarray:
"""
Horizontal binning of the image Z

Parameters
----------
Z : np.ndarray
The image to be binned.
k : int
Number of iterations to bin the image by half.

fac : int
binning factor
dim : direction X=0, Y=1
Returns
-------
list[np.ndarray]
List of k images.
np.ndarray
binned image

Example
-------
>>> Z = np.random.rand(64, 64)
>>> binned_zs = horizontal_binning(Z, 3)
>>> len(binned_zs)
4
"""
binned_zs = [Z]
for _ in range(k):
sub_z0 = Z[:, ::2]
sub_z1 = Z[:, 1::2]
# make sure z0 and z1 have the same shape
if sub_z0.shape[1] > sub_z1.shape[1]:
sub_z0 = sub_z0[:, :-1]
elif sub_z0.shape[1] < sub_z1.shape[1]:
sub_z1 = sub_z1[:, :-1]
# average z0 and z1
Z = (sub_z0 + sub_z1) * 0.5
binned_zs.append(Z)
return binned_zs


def horizontal_debinning(original_image: np.ndarray, target: np.ndarray) -> np.ndarray:

if fac > 1:
fac_half = fac // 2
binned_zs = create_array(Z, fac, dim)

# get coordinates of bin centres
if dim == 0:
binned_zs = binned_zs[
fac_half + ((fac % 2) == 1) : binned_zs.shape[dim] - fac_half + 1 : fac,
:,
]
else:
binned_zs = binned_zs[
:,
fac_half + ((fac % 2) == 1) : binned_zs.shape[dim] - fac_half + 1 : fac,
]

return binned_zs

return Z


def horizontal_debinning(
Z: np.ndarray, size: int, fac: int, n_iter: int, dim: int = 1
) -> np.ndarray:
"""
Horizontal debinning of the image Z into the same shape as Z_target.

Parameters
----------
original_image : np.ndarray
Z : np.ndarray
The image to be debinned.
target : np.ndarray
The target image to match the shape.
size: target size (original size before binning) for the second dimension
fac: binning factor (original divisor)
n_iter: number of iterations
dim: dimension for binning (Y = 0 or X = 1)

Returns
-------
np.ndarray
The debinned image.

Example
-------
>>> Z = np.random.rand(64, 64)
>>> target = np.random.rand(64, 128)
>>> debinned_z = horizontal_debinning(Z, target)
>>> debinned_z.shape
(64, 128)
"""
# Original dimensions
original_height, original_width = original_image.shape
# Target dimensions
new_height, new_width = target.shape

# Original grid
original_x = np.arange(original_width)
original_y = np.arange(original_height)

# Target grid
new_x = np.linspace(0, original_width - 1, new_width)
new_y = np.linspace(0, original_height - 1, new_height)

# Spline interpolation
spline = RectBivariateSpline(original_y, original_x, original_image)
interpolated_image = spline(new_y, new_x)
if fac <= 1:
return np.copy(Z)

fac_half = fac // 2

if dim == 0:
base_array = np.ones((size, 1))
else:
base_array = np.ones((1, size))

n_counter = create_array(base_array, fac, dim)

# coordinates of bin counts
x1c = np.arange(fac_half + ((fac % 2) == 1), (Z.shape[dim]) * fac, fac)
x1 = np.arange(fac_half + 1 - ((fac % 2) == 0) / 2, (Z.shape[dim]) * fac, fac)

# coordinates of image pixels
ix1 = np.arange(1, size + 1)

interpolated_image = 0

for j in range(max(1, n_iter)):
# residual
if j > 0:
residual = Z - horizontal_binning(interpolated_image, fac, dim)
else:
residual = Z

# interpolation
if dim == 0:
interp = interp1d(
x1,
residual / n_counter[x1c, :],
kind="cubic",
fill_value="extrapolate",
axis=0,
)
else:
interp = interp1d(
x1, residual / n_counter[:, x1c], kind="cubic", fill_value="extrapolate"
)
interpolated_image = interpolated_image + interp(ix1)

return interpolated_image

Expand Down
49 changes: 16 additions & 33 deletions tests/unit/bm3dornl/test_bm3d_ring_artifact_removal_ms.py
Original file line number Diff line number Diff line change
@@ -1,53 +1,36 @@
import pytest
from unittest.mock import patch
import numpy as np
from bm3dornl.bm3d import bm3d_ring_artifact_removal_ms


size_x = 256
size_y = 256


@pytest.fixture
def setup_sinogram():
return np.random.rand(256, 256)
return np.random.rand(size_x, size_y)


@patch("bm3dornl.bm3d.horizontal_binning")
@patch("bm3dornl.bm3d.horizontal_debinning")
@patch("bm3dornl.bm3d.bm3d_ring_artifact_removal")
@pytest.mark.cuda_required
def test_bm3d_ring_artifact_removal_ms(
mock_bm3d_ring_artifact_removal,
mock_horizontal_debinning,
mock_horizontal_binning,
setup_sinogram,
):
sinogram = setup_sinogram
binned_sinos = [np.random.rand(64, 64) for _ in range(4)]
mock_horizontal_binning.return_value = binned_sinos
binned_sinos = [np.random.rand(64, 64) for _ in range(4)]
mock_bm3d_ring_artifact_removal.return_value = binned_sinos
binned_sinos = [np.random.rand(64, 64) for _ in range(4)]
mock_horizontal_debinning.return_value = binned_sinos

result = bm3d_ring_artifact_removal_ms(sinogram, k=4)

assert result is not None
mock_horizontal_binning.assert_called_once_with(sinogram, k=4)
assert mock_bm3d_ring_artifact_removal.call_count == 3
assert mock_horizontal_debinning.call_count == 3

r, c = result.shape

assert c == size_x
assert r == size_y

result_single_pass = bm3d_ring_artifact_removal_ms(sinogram, k=0)
assert result_single_pass is not None
mock_bm3d_ring_artifact_removal.assert_called_with(
sinogram,
mode="simple",
block_matching_kwargs={
"patch_size": (8, 8),
"stride": 3,
"background_threshold": 0.0,
"cut_off_distance": (64, 64),
"num_patches_per_group": 32,
"padding_mode": "circular",
},
filter_kwargs={
"filter_function": "fft",
"shrinkage_factor": 3e-2,
},
)

r, c = result_single_pass.shape

assert c == size_x
assert r == size_y
Loading
Loading