comp-imaging / ProxImaL

A domain-specific language for image optimization.
MIT License
112 stars 29 forks source link

New example script: 2D deconvolution with spatially-variant blur kernel #84

Closed shnaqvi closed 10 months ago

shnaqvi commented 1 year ago

for testing the spatially varying PSF deconvolution algorithm from Flicker & Rigaugt 2005

antonysigma commented 1 year ago

Thank for submitting the draft script. I will test it on my computer tomorrow.

@shnaqvi I tested the script by setting a smaller image size: im = np.random.randn(128, 128). The solver completes within ~2 iterations~ with timing log.

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

No, mine didn't return any ValueError.

antonysigma commented 1 year ago

Actually no, I am getting the output from the blur forward model with this command and I want to define the prior on this output.

@shnaqvi I am curious about the work that advocates for a prior on the blurred image, i.e. nonneg(blurred) + norm1(grad(blurred)). May I have the citation and the DOI hyperlink to the article?

shnaqvi commented 1 year ago

Thanks @antonysigma. Thanks for fixing it. Not sure, what was the issue that reshape fixed. I got another error related to numpy compatibility, which you might consider fixing for future compliance. I was able to unblock myself by just 1 line of code change in comp_graph.py return statement:

Traceback (most recent call last):
  File "examples/test_deconv_sv_psf.py", line 75, in <module>
    prob.solve(solver='pc', max_iters=2, verbose=True, x0=im.copy())
  File "../../../../../.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/algorithms/problem.py", line 156, in solve
    Knorm = est_CompGraph_norm(K, try_fast_norm=self.try_fast_norm)
  File "../../../../../.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/lin_ops/comp_graph.py", line 433, in est_CompGraph_norm
    return np.float(Knorm)
  File "../../../../../.pyenv/versions/3.10.1/lib/python3.10/site-packages/numpy/__init__.py", line 305, in __getattr__
    raise AttributeError(__former_attrs__[attr])
AttributeError: module 'numpy' has no attribute 'float'.
`np.float` was a deprecated alias for the builtin `float`. To avoid this error in existing code, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.
The aliases was originally deprecated in NumPy 1.20; for more details and guidance see the original release note at:
    https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations

Regarding adding priors, I can refer you to this Nature paper that recently used this technique in microscopy and used L1 prior during optimization. https://doi.org/10.1038/s41377-020-00403-7

shnaqvi commented 1 year ago

Yes we do except PRs to resolve the np.float vs float issues. I will be on standby, waiting for your PR. Or, I can submit the one-liner fix myself. Either way is fine for me.

I think it'd be okay if you push that change as I'm still struggling with getting my example to work.

shnaqvi commented 1 year ago

Thanks. The paper does require nonneg(x) and norm1(grad(x)) to "inverse" the forward problem though. No worries, it doesn't impact the scope and the PR. ProxImaL is robust enough for users (i.e. you) to lead the problem formulations.

Sorry I didn't understand your comment fully. Can you elaborate on what you mean by "to inverse the forward problem." This is what we are doing using proximal right, solving the inverse problem iteratively? Are you suggesting that the paper is doing something else?

antonysigma commented 1 year ago

Yes we do accept PRs to resolve the np.float vs float issues. I will be on standby, waiting for your PR. Or, I can submit the one-liner fix myself. Either way is fine for me.

I think it'd be okay if you push that change as I'm still struggling with getting my example to work.

Discussion moved to a separate issue #85 .

shnaqvi commented 1 year ago

Sorry for the delay in replying @antonysigma. I actually tried the code on my data but even though the solver runs now, it does not return even after an hour for just 2 iterations.

So I went ahead and tried running it on a simpler problem of Spatially "Invariant" PSF deblurring using synthetic data with ground truth. I've added this as a new script test_deconv_sinv_psf.py. Here, the solver runs and also returns with 2 iterations in 2 mins with blurred image used as the initial condition. However, I'm seeing these issues:

  1. recovered image is exactly the same as the blurred image with no increase in contrast
  2. mu parameter for the TV operator, which should control the blockiness of the image seems to have no effect, even if it's set to a really large number of 1e3.

I believe I've setup the adjoint of the forward model correctly as the adjoint of a simple convolution forward model with a kernel will be a correlation (or convolution with the rotated kernel).

Can you please look at this code and comment? I think if we make this simple example to work then we can fix the issue with the original Spatially "Varying" PSF problem at hand.

shnaqvi commented 1 year ago
image
shnaqvi commented 1 year ago
m_grad = px.grad(x, implem='halide')
grad_term = mu * px.group_norm1(im_grad, group_dims=(2, ), implem='halide')

@antonysigma, with your changes to the test_deconv_sinv_psf.py in commit 33cb503, I'm getting proximal.halide.build Module not found error. Do I need to install something for it to accept implem='halide' flag?:

Traceback (most recent call last):
File "/Users/salman_naqvi/Documents/Programming/VisualStudio/ProxImaL/proximal/examples/test_deconv_sinv_psf.py", line 86, in <module>
prob.solve(solver='ladmm', eps_abs=1e-2, max_iters=20, verbose=True) #eps_abs=1e-2, 
File "/Users/salman_naqvi/.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/algorithms/problem.py", line 199, in solve
opt_val = module.solve(psi_fns, omega_fns,
File "/Users/salman_naqvi/.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/algorithms/linearized_admm.py", line 94, in solve
K.forward(v, Kv)
File "/Users/salman_naqvi/.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/lin_ops/comp_graph.py", line 271, in forward
self.traverse_graph(forward_eval, True)
File "/Users/salman_naqvi/.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/lin_ops/comp_graph.py", line 328, in traverse_graph
node_fn(curr)
File "/Users/salman_naqvi/.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/lin_ops/comp_graph.py", line 263, in forward_eval
node.forward(inputs, outputs)
File "/Users/salman_naqvi/.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/lin_ops/grad.py", line 51, in forward
Halide('A_grad').A_grad(tmpin, self.tmpfwd)  # Call
File "/Users/salman_naqvi/.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/halide/halide.py", line 63, in <lambda>
setattr(self, self.func, lambda *args: self.run(*args))
File "/Users/salman_naqvi/.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/halide/halide.py", line 96, in run
launch = importlib.import_module(
File "/Users/salman_naqvi/.pyenv/versions/3.10.1/lib/python3.10/importlib/__init__.py", line 126, in import_module
return _bootstrap._gcd_import(name[level:], package, level)
File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
File "<frozen importlib._bootstrap>", line 992, in _find_and_load_unlocked
File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
File "<frozen importlib._bootstrap>", line 1004, in _find_and_load_unlocked
ModuleNotFoundError: No module named 'proximal.halide.build'
antonysigma commented 1 year ago

Do I need to install something for it to accept implem='halide' flag?:

@shnaqvi Looks like the Halide-accelerated computing module is disabled in your system. Please remove the code ..., implem='halide' and then synchronize the changes to this PR. The feature so far is irrelevant to your problem formulation needs.

The halide-python bindings and the design intent is described in the article https://stevendiamond.me/pdf/proximal.pdf . It aims to achieve ~5X compute speed by eliminating the data movement cost in Numpy/Scipy. But as a trade off, halide module needs direct access to the C++ compiler and (optionally) the metal/cuda gpu hardware in your MacOS environment. As I recalled, end users are unwilling to grant ProxImaL direct access to the native system. Therefore, we do not have enough volunteers and machines to QC the module.

shnaqvi commented 1 year ago
# Compute the norm of the Gram matrix, i.e. max || convT(conv(x)) || for all images x.
 scaling_factor = np.sum(psf)
 return LinOpFactory(im_shape, im_shape, forward, adjoint, norm_bound=scaling_factor)

@antonysigma , FYI, adding the norm_bound to the BlackBox implementation of the forward model makes all the difference, and now it is at par with the px.conv() operator, returning in 2 mins when psf is not cropped! Without this flag, it wouldn't return in 10s of mins. And if the cropped PSF is used then it returns in just under a minute. But now its even faster with pc Chambolle-Pock and cropped PSF, returning in half a minute. 🙂.

So now the question is, how do I find the norm_bound in case of the more complex forward model for the Spatially Variant PSF case? There, I'm summing up the weighted convolutions, so would it be the sum over sum over all the PSFs?

shnaqvi commented 1 year ago

Sorry @antonysigma for missing in action. I pulled your changes and tried running the test_deconv_sv_psf.py script but its giving me a ValueError about array having inhomogeneous shape, which is being raised inside proximal/algorithms/problem.py during a list comprehension on prox_fns. I see that you have changed the logic of the forward and adjoint functions and instead of summing the convolutions in the forward function, you are appending them to a list and then adding inside the adjoint function. Since I didn't understand this change, so thought of asking you if this could be breaking the code? I tried debugging but haven't been able to resolve it. Were you able to run it successfully on your end?

Traceback (most recent call last):
  File "examples/test_deconv_sv_psf.py", line 171, in <module>
    prob.solve(solver='pc', eps_abs=1e-2, max_iter=maxiter, verbose=True, x0=im)
  File "../../../../../.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/algorithms/problem.py", line 108, in solve
    prox_fns = [absorb.absorb_offset(fn) for fn in prox_fns]
  File "../../../../../.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/algorithms/problem.py", line 108, in <listcomp>
    prox_fns = [absorb.absorb_offset(fn) for fn in prox_fns]
  File "../../../../../.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/algorithms/absorb.py", line 111, in absorb_offset
    new_b = -prox_fn.lin_op.get_offset()
  File "../../../../../.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/lin_ops/lin_op.py", line 223, in get_offset
    offset = self.value
  File "../../../../../.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/lin_ops/lin_op.py", line 213, in value
    self.forward(inputs, [output])
  File "../../../../../.pyenv/versions/3.10.1/lib/python3.10/site-packages/proximal/lin_ops/sum.py", line 20, in forward
    np.copyto(outputs[0], np.sum(inputs, 0))
  File "<__array_function__ internals>", line 200, in sum
  File "../../../../../.pyenv/versions/3.10.1/lib/python3.10/site-packages/numpy/core/fromnumeric.py", line 2324, in sum
    return _wrapreduction(a, np.add, 'sum', axis, dtype, out, keepdims=keepdims,
  File "../../../../../.pyenv/versions/3.10.1/lib/python3.10/site-packages/numpy/core/fromnumeric.py", line 86, in _wrapreduction
    return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 1 dimensions. The detected shape was (2,) + inhomogeneous part.
shnaqvi commented 1 year ago

Also regarding the space-invariant case, I generated results using with px.conv() and then using BlackBox based convolution forward operator by toggling the using_px_conv flag, and the results are visibly bad with the latter. Do you know why?

image