Skip to content

Commit

Permalink
Example script clean up (lanl#50)
Browse files Browse the repository at this point in the history
* Change test image

* Add missing doi and various other clean up

* Edit header docstring

* Add missing month fields
  • Loading branch information
bwohlberg authored Oct 21, 2021
1 parent 9357bb1 commit 76db0a8
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 72 deletions.
29 changes: 18 additions & 11 deletions docs/source/references.bib
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ @Article {almeida-2013-deconvolving
title = {Deconvolving Images With Unknown Boundaries Using
the Alternating Direction Method of Multipliers},
year = 2013,
month = Aug,
volume = 22,
number = 8,
pages = {3074--3086},
Expand All @@ -16,7 +17,9 @@ @Article {barzilai-1988-stepsize
journal = {{IMA} Journal of Numerical Analysis},
volume = 8,
pages = {141--148},
year = 1988
year = 1988,
month = Jan,
doi = {10.1093/imanum/8.1.141}
}

@Article {beck-2009-fast,
Expand All @@ -37,9 +40,11 @@ @Article {beck-2009-tv
author = {Beck, Amir and Teboulle, Marc},
journal = {IEEE Transactions on Image Processing},
year = 2009,
month = Nov,
volume = 18,
number = 11,
pages = {2419--2434}
pages = {2419--2434},
doi = {10.1109/TIP.2009.2028250}
}

@Book {beck-2017-first,
Expand Down Expand Up @@ -88,8 +93,8 @@ @Article {clinthorne-1993-preconditioning
volume = 12,
number = 1,
pages = {78--83},
month = mar,
doi = {10.1109/42.222670},
month = Mar,
doi = {10.1109/42.222670}
}

@InProceedings{dabov-2008-image,
Expand All @@ -105,6 +110,7 @@ @InProceedings{dabov-2008-image
publisher = {SPIE},
pages = {62--73},
year = 2008,
month = Mar,
doi = {10.1117/12.766355}
}

Expand All @@ -114,6 +120,7 @@ @InProceedings {florea-2017-robust
booktitle = {Proceedings of the IEEE International Conference on
Acoustics, Speech and Signal Processing (ICASSP)},
year = 2017,
month = Mar,
pages = {4521--4525},
doi = {10.1109/ICASSP.2017.7953012},
location = {New Orleans, LA, USA}
Expand Down Expand Up @@ -148,14 +155,13 @@ @Article {glowinski-1975-approximation
url = {http://eudml.org/doc/193269}
}

@Article {goldstein-2014-fasta,
@Misc {goldstein-2014-fasta,
title = {A Field Guide to Forward-Backward Splitting with a
FASTA Implementation},
{FASTA} Implementation},
author = {Tom Goldstein and Christoph Studer and Richard
Baraniuk},
year = 2014,
journal = {arXiv eprint},
volume = {abs/1411.3406},
eprint = {arXiv:1411.3406},
url = {http://arxiv.org/abs/1411.3406},
}

Expand Down Expand Up @@ -216,7 +222,7 @@ @Article {menon-2007-demosaicing
Calvagno},
journal = {IEEE Transactions on Image Processing},
year = 2007,
month = jan,
month = Jan,
volume = 16,
number = 1,
pages = {132--141},
Expand All @@ -237,13 +243,14 @@ @Misc {pfister-2021-scico
@Article {sauer-1993-local,
title = {A local update strategy for iterative reconstruction
from projections},
author = {Sauer, K. and Bouman, C.},
author = {Sauer, Ken and Bouman, Charles},
journal = {IEEE Transactions on Signal Processing},
year = 1993,
month = Feb,
number = 2,
pages = {534--548},
volume = 41,
doi = {10.1109/78.193196},
doi = {10.1109/78.193196}
}

@Article {sreehari-2016-plug,
Expand Down
49 changes: 20 additions & 29 deletions examples/scripts/admm_tv_isotropic.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,66 +5,56 @@
# with the package.

r"""
Isotropic Total Variation
=========================
Isotropic Total Variation (ADMM)
================================
This example demonstrates isotropic total variation (TV)
regularization. It solves the denoising problem
This example compares denoising via isotropic and anisotropic total
variation (TV) regularization. It solves the denoising problem
$$\mathrm{argmin}_{\mathbf{x}} \; (1/2) \| \mathbf{y} - \mathbf{x}
\|^2 + \lambda R(\mathbf{x}) \;,$$
where $R$ is the isotropic TV: the sum of the norms of the gradient
vectors at each point in the image $\mathbf{x}$. The same
reconstruction is performed with anisotropic TV regularization for
comparison; the isotropic version shows fewer block-like artifacts.
where $R$ is either the isotropic or anisotropic TV regularizer.
In SCICO, switching between these two regularizers is a one-line
change: replacing an
[L1Norm](../_autosummary/scico.functional.rst#scico.functional.L1Norm)
with a
[L21Norm](../_autosummary/scico.functional.rst#scico.functional.L21Norm).
Note that the isotropic version exhibits fewer block-like artifacts on
edges that are not vertical or horizontal.
"""

import jax
import jax.numpy as jnp
import jax.scipy as jsp

from xdesign import SiemensStar, discrete_phantom

import scico.numpy as snp
import scico.random
from scico import functional, linop, loss, plot
from scico.admm import ADMM, LinearSubproblemSolver

"""
Create a ground truth image.
"""
N = 256 # image size

# Create a ground truth image by spatially filtering noise.
kernel_size = N // 5
key = jax.random.PRNGKey(1)
x_gt = jax.random.uniform(key, shape=(N + kernel_size - 1, N + kernel_size - 1))
x = jnp.linspace(-3, 3, kernel_size)
window = jsp.stats.norm.pdf(x) * jsp.stats.norm.pdf(x[:, None])
window = window / window.sum()
x_gt = jsp.signal.convolve(x_gt, window, mode="valid")
x_gt = (x_gt > jnp.percentile(x_gt, 25)).astype(float) + (x_gt > jnp.percentile(x_gt, 75)).astype(
float
)
phantom = SiemensStar(16)
x_gt = snp.pad(discrete_phantom(phantom, 240), 8)
x_gt = jax.device_put(x_gt) # convert to jax type, push to GPU
x_gt = x_gt / x_gt.max()


"""
Add noise to create a noisy test image.
"""
σ = 1.0 # noise standard deviation
key, subkey = jax.random.split(key)
n = σ * jax.random.normal(subkey, shape=x_gt.shape)
y = x_gt + n
σ = 0.75 # noise standard deviation
noise, key = scico.random.randn(x_gt.shape, seed=0)
y = x_gt + σ * noise


"""
Denoise with isotropic total variation
"""
reg_weight_iso = 2e0
reg_weight_iso = 1.4e0
f = loss.SquaredL2Loss(y=y)
g_iso = reg_weight_iso * functional.L21Norm()

Expand All @@ -91,7 +81,7 @@
Denoise with anisotropic total variation for comparison.
"""
# Tune the weight to give the same data fidelty as the isotropic case.
reg_weight_aniso = 1.74e0
reg_weight_aniso = 1.2e0
g_aniso = reg_weight_aniso * functional.L1Norm()

solver = ADMM(
Expand Down Expand Up @@ -149,4 +139,5 @@
fig.suptitle("Denoising comparison (zoomed)")
fig.show()


input("\nWaiting for input to close figures and exit")
60 changes: 28 additions & 32 deletions examples/scripts/pgm_tv_isotropic.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,27 +14,30 @@
denoising problem
$$\mathrm{argmin}_{\mathbf{x}} \; (1/2) \| \mathbf{y} - \mathbf{x}
\|^2 + \lambda R(\mathbf{x}) \;,$$
\|^2 + \lambda R(\mathbf{x}) + \iota_C(\mathbf{x}) \;,$$
where $R$ is the isotropic TV: the sum of the norms of the gradient
vectors at each point in the image $\mathbf{x}$. The same
reconstruction is performed with anisotropic TV regularization for
comparison; the isotropic version shows fewer block-like artifacts.
where $R$ is a TV regularizer, $\iota_C(\cdot)$ is the indicator function
of constraint set $C$, and $C = \{ \mathbf{x} \, | \, x_i \in [0, 1] \}$,
i.e. the set of vectors with components constrained to be in the interval
$[0, 1]$. The problem is solved seperately with $R$ taken as isotropic
and anisotropic TV regularization
The solution via PGM is based on :cite:`beck-2009-tv`. This follows a
dual approach that constructs a dual for the constrained denoising
problem (the constraint given by restricting the solution to the [0,1]
range). The PGM solution minimizes the resulting dual. In this case,
switching between the two regularizers corresponds to switching
between two different projectors.
The solution via PGM is based on the approach in :cite:`beck-2009-tv`,
which involves constructing a dual for the constrained denoising problem.
The PGM solution minimizes the resulting dual. In this case, switching
between the two regularizers corresponds to switching between two
different projectors.
"""

from typing import Callable, Optional, Union

import jax
import jax.numpy as jnp
import jax.scipy as jsp

from xdesign import SiemensStar, discrete_phantom

import scico.numpy as snp
import scico.random
from scico import functional, linop, loss, operator, plot
from scico.blockarray import BlockArray
from scico.pgm import AcceleratedPGM, RobustLineSearchStepSize
Expand All @@ -45,29 +48,18 @@
Create a ground truth image.
"""
N = 256 # image size


# Create a ground truth image by spatially filtering noise.
kernel_size = N // 5
key = jax.random.PRNGKey(1)
x_gt = jax.random.uniform(key, shape=(N + kernel_size - 1, N + kernel_size - 1))
x = jnp.linspace(-3, 3, kernel_size)
window = jsp.stats.norm.pdf(x) * jsp.stats.norm.pdf(x[:, None])
window = window / window.sum()
x_gt = jsp.signal.convolve(x_gt, window, mode="valid")
x_gt = (x_gt > jnp.percentile(x_gt, 25)).astype(float) + (x_gt > jnp.percentile(x_gt, 75)).astype(
float
)
phantom = SiemensStar(16)
x_gt = snp.pad(discrete_phantom(phantom, 240), 8)
x_gt = jax.device_put(x_gt) # convert to jax type, push to GPU
x_gt = x_gt / x_gt.max()


"""
Add noise to create a noisy test image.
"""
σ = 1.0 # noise standard deviation
key, subkey = jax.random.split(key)
n = σ * jax.random.normal(subkey, shape=x_gt.shape)
y = x_gt + n
σ = 0.75 # noise standard deviation
noise, key = scico.random.randn(x_gt.shape, seed=0)
y = x_gt + σ * noise


"""
Expand Down Expand Up @@ -113,6 +105,8 @@ def __call__(self, x: Union[JaxArray, BlockArray]) -> float:
Denoise with isotropic total variation. Define projector for isotropic
total variation.
"""


# Evaluation of functional set to zero.
class IsoProjector(functional.Functional):

Expand All @@ -139,7 +133,7 @@ def prox(self, x: JaxArray, lam: float) -> JaxArray:
Use RobustLineSearchStepSize object and set up AcceleratedPGM solver
object. Run the solver.
"""
reg_weight_iso = 2e0
reg_weight_iso = 1.4e0
f_iso = DualTVLoss(y=y, A=A, lmbda=reg_weight_iso)
f_iso.is_smooth = True
g_iso = IsoProjector()
Expand All @@ -165,6 +159,8 @@ def prox(self, x: JaxArray, lam: float) -> JaxArray:
Denoise with anisotropic total variation for comparison. Define
projector for anisotropic total variation.
"""


# Evaluation of functional set to zero.
class AnisoProjector(functional.Functional):

Expand All @@ -186,7 +182,7 @@ def prox(self, x: JaxArray, lam: float) -> JaxArray:
isotropic case. Run the solver.
"""

reg_weight_aniso = 1.74e0
reg_weight_aniso = 1.2e0
f = DualTVLoss(y=y, A=A, lmbda=reg_weight_aniso)
f.is_smooth = True
g = AnisoProjector()
Expand Down Expand Up @@ -233,7 +229,7 @@ def prox(self, x: JaxArray, lam: float) -> JaxArray:
fig.suptitle("Denoising comparison")
fig.show()

# Zoomed version
# zoomed version
fig, ax = plot.subplots(nrows=2, ncols=2, sharex=True, sharey=True, figsize=(11, 10))
plot.imview(x_gt, title="Ground truth", fig=fig, ax=ax[0, 0], **plt_args)
plot.imview(y, title="Noisy version", fig=fig, ax=ax[0, 1], **plt_args)
Expand Down

0 comments on commit 76db0a8

Please sign in to comment.