An image reconstruction method based on a differentiable proximal algorithm

Abstract

This invention relates to an image reconstruction method based on a differentiable proximal algorithm, belonging to the field of computer vision technology. First, a proximal algorithm model for image reconstruction is performed. Then, a proximal algorithm programming framework for image reconstruction is designed. Next, the differential design of the proximal algorithm is performed. Finally, the image reconstruction is performed using the proximal algorithm and framework. This invention introduces a domain-specific modeling language and compiler to solve the memory explosion and operational difficulties problems in image reconstruction methods using differentiable proximal algorithms. Only a few lines of code are needed to generate high-performance solvers for a series of image reconstruction problems, including end-to-end image super-resolution, image deraining, image denoising, and image deblurring. Compared with existing technologies, this invention achieves better image reconstruction results, significantly reduces memory requirements, has strong portability, is simple to use, and can be applied to more application scenarios.

Description

Technical Field

[0001] This invention relates to an image reconstruction method based on a differentiable proximal algorithm, belonging to the field of computer vision technology. Background Technology

[0002] Computer vision is a simulation of biological vision using computers and related equipment. Its main task is to process acquired images or videos to obtain three-dimensional information of the corresponding scene.

[0003] Image reconstruction, aiming to recover a clean underlying image from low-level observations, is an important task in the field of computer vision. The image reconstruction problem is a typical inverse problem. In practice, because the inverse mapping is unknown, the solution space for image reconstruction is infinite, requiring the application of regularization techniques to obtain feasible and optimal solutions. Therefore, most research on image restoration focuses on employing effective analytical models and learning schemes to find approximations of the exact mapping to recover degraded images.

[0004] A series of studies have shown that combining proximal optimization-based methods with some learning-based modules can create a highly effective and interpretable hybrid approach for image reconstruction, allowing image reconstruction to be performed with almost no additional training data. However, manually implementing and testing this hybrid approach for different image reconstruction tasks requires domain expertise in both proximal optimization and deep learning, which is often error-prone and very time-consuming.

[0005] Furthermore, simply expanding these iterative optimization methods to generate a computational graph that can be differentiated using automatic differentiation techniques in order to train learning-based modules leads to an explosive increase in memory consumption, making batch-based training extremely difficult. In addition, some image reconstruction tasks present optimization challenges for non-differentiable paths, preventing the use of differentiable proximal algorithms.

[0006] Currently, no publicly available image reconstruction method based on differentiable proximal algorithms has been found that meets practical needs. Summary of the Invention

[0007] The purpose of this invention is to address the technical problems of image reconstruction methods, such as the difficulty in constructing algorithms and poor practicality. To overcome the shortcomings of existing technologies, this invention creatively proposes an image reconstruction method based on a differentiable proximal-end algorithm. This invention proposes a series of memory optimization strategies and designs a programming framework to reduce programming difficulty, thereby realizing image reconstruction based on the differentiable proximal-end algorithm.

[0008] This invention introduces It is a domain-specific modeling language and compiler that addresses the memory explosion and operational difficulties in image reconstruction methods using differentiable proximal algorithms. It allows users to concisely specify optimization objective functions containing unknown variables at a high level and intelligently compiles the objective functions into a highly computationally and memory-efficient differentiable solver. One of its core features is its complete differentiability, which supports a hybrid of form-based optimization and learning-based solvers. This integrates near-end optimization with neural network solving. It supports a wide range of applications, including learning-based prior optimization objective functions and sample-related inner-loop optimization schedulers built based on deep equilibrium learning or deep reinforcement learning. Therefore, only a few lines of code are needed. It can generate high-performance solvers for a range of image reconstruction problems, including end-to-end image super-resolution, image deraining, image denoising, and image deblurring.

[0009] First, we model the proximal algorithm for image reconstruction. Then, we design a proximal algorithm programming framework for image reconstruction. Next, we perform the differential design of the proximal algorithm for image reconstruction. Finally, we use the proximal algorithm and framework to perform image reconstruction.

[0010] Beneficial effects

[0011] Compared with the prior art, the present invention has the following advantages:

[0012] 1. Improved image reconstruction results. Compared to existing technologies, the results of this invention outperform existing methods across various evaluation metrics.

[0013] 2. Significantly reduced video memory requirements. Compared to existing technologies, this invention requires less video memory and is suitable for a wider range of applications.

[0014] 3. Highly transferable and easy to use. Compared with existing technologies, this method is simple and easy to use, and can be easily transferred to different image reconstruction tasks. Attached Figure Description

[0015] Figure 1 This is a schematic diagram of the core algorithm model of the method of this invention.

[0016] Figure 2 This is a schematic diagram of the system composition of the present invention. Specific implementation methods

[0017] The present invention will now be described in further detail with reference to the accompanying drawings.

[0018] Specifically, the present invention is implemented using the following technical solutions.

[0019] An image reconstruction method based on a differentiable proximal-end algorithm includes the following steps:

[0020] Step 1: Modeling the near-end algorithm for image reconstruction.

[0021] Proximal algorithms, as flexible optimization algorithms, have been widely adopted in various image reconstruction fields in recent years. Different algorithm variants and applications are usually proposed and implemented independently in each work. To achieve optimal performance, each individual image reconstruction work requires tedious experimentation with various proximal algorithms, proximal operators, and constraint terms. This greatly limits the further development and utilization of proximal algorithms.

[0022] Therefore, this invention proposes for the first time a unified proximal algorithm modeling method for image reconstruction. All optimization problems considered in this invention can be transformed into a proximal canonical form for solution.

[0023] Specifically, the objective of each standard form solution is to find an unknown variable by minimizing an objective function. Let f denote the real number field, and n denote the vector dimension. This objective function consists of a series of penalty functions f. i The sum is composed of, and each penalty function takes an unknown variable after a certain linear transformation as input K. i x. In addition, the standard form has a series of additional constraints c. j As shown below:

[0024]

[0025]

[0026] Where argmin represents the value of the variable that minimizes the latter expression; st indicates compliance; I represents the number of penalty functions; θ represents the implicit parameter; c represents the constraint condition; and f is the penalty function. J represents any; J represents the number of constraints; K represents any. i For linear operators, Responsible for projecting the unknowns onto the penalty function On the input, constraints Add m by means of equality or inequality. i This indicates the number of rows in the matrix.

[0027] Here, f i y i and x j Through a set of implicit parameters Parameterization is performed. For learnable latent parameters, the proximal canonical form assumes that the linear operator is differentiable in the sense of machine learning (e.g., neural networks or differentiable physical forward models). Once the proximal canonical form is compiled into a differentiable proximal solver, this canonical form can achieve backpropagation, thereby using data-driven end-to-end learning to optimize the latent parameters and the differentiable internal solver parameters.

[0028] Step 2: Design a proximal algorithm programming framework for image reconstruction.

[0029] This invention designs an intuitive, concise, and powerful domain modeling language as a programming interface for modeling differentiable proximal algorithms for image reconstruction. This interface allows users to concisely specify objective optimization functions based on specific image reconstruction problems at a high level, and integrates various built-in and user-defined modeling components, such as various proximal algorithms, proximal operators, linear operators, hyperparameter generators, etc.

[0030] To define a specific proximal optimization problem, this invention introduces It is a domain-specific modeling language and compiler that addresses the memory explosion and operational inefficiency problems in image reconstruction methods using differentiable proximal algorithms. First, define an unknown variable x and a placeholder representing the observed input y. Then, define each term in the optimization objective by combining linear operators (LinOp) and proximal functions / operators (ProxFn) that may have latent parameters. The linear operators (LinOp) can be arbitrarily combined to form a linear expression, where each leaf node is a variable. Each proximal function (ProxFn) accepts a linear expression and optional keyword arguments as input, based on its specific function signature. After defining all the proximal functions (ProxFn), combine them into a single optimization objective and wrap it in a problem container to construct the corresponding optimization problem.

[0031] The following describes the core primitives and their role in... Functions in language.

[0032] Variable: refers to a multidimensional array or a one-dimensional vector.

[0033] It is not required to specify the shape of each variable in advance. A built-in mechanism is provided to automatically infer the shape of each variable before solving the problem and propagate it to each LinOp and ProxFn for necessary shape-related initialization. Unlike variables that are bound to a given input shape (where the user must create a new problem for inputs with different shapes), this feature not only reduces the user's burden but also allows for optimization of different input shapes in multi-task training.

[0034] Placeholder: A placeholder that represents the actual data.

[0035] Placeholders allow a problem to be reused multiple times when it needs to address different input observations simultaneously. This is common when evaluating an entire dataset of observations. For example, when using a placeholder y, a user can freely modify the underlying data of y after creating the problem by assigning a value (i.e., y.value = input), thus avoiding redefining the problem.

[0036] Parameters: A key primitive for enabling differentiable optimization is that it encodes learnable parameters associated with differentiable operators (such as LinOp and ProxFn). Params provide the ability to make the behavior of operators adaptive to the problem when given a higher-level dataset for optimizing the entire proximal solver.

[0037] Params can be explicitly defined by the user (e.g., manually constructing a denoising neural network as a deep prior) or implicitly defined within the operator itself.

[0038] Linear operators (LinOp): responsible for processing data from multidimensional space Mapped to another multidimensional space

[0039] All linear operators are differentiable, and evaluation of the forward operation x→Kx and the adjoint operation x→K is allowed. T x, where K represents and T represents. It supports combining any number of existing linear operators and automatically converts them into directed acyclic graphs (DAGs) for evaluating their positive and adjoint operations.

[0040] ProxFn: Each penalty function in the optimization objective has a corresponding proximal operator, called a proximal function. Each near-endian function in the class inherits from the base class ProxFn and implements fn.eval() and fn.prrox() for evaluating the function and its near-endian operators. If the function cannot be explicitly evaluated, fn.eval() is omitted. All proximal operators are differentiable and can be effectively used as basic components to build optimal, matrix-free, differentiable solvers.

[0041] The following shows an applicable example. The typical proximal optimization problem described is:

[0042] x = Variable()

[0043] y = Placeholder(input)

[0044] data_term=sum_squares(conv(x,psf_DOE)-y)

[0045] prior_term=deep_prior(x,unet)

[0046] objective=data_term+prior_term+nonneg(x)

[0047] p = Problem(objective)

[0048] This includes different basic modules such as ProxFn, LinOp, and Params.

[0049] Step 3: Differential design of the near-end algorithm for image reconstruction.

[0050] Deep learning-based learnable methods have achieved remarkable success and development in the field of image reconstruction in recent years. Some works have shown that combining formal proximal algorithms with learning-based algorithms can yield a hybrid method that is more powerful and more interpretable than deep learning methods.

[0051] To effectively support the modeling of near-end algorithms, this invention proposes three differentiation strategies, including iterative expansion, equilibrium learning, and reinforcement learning. Details are as follows:

[0052] Iterative expansion: Support for a differential strategy based on iterative expansion. This strategy expands the iterative optimization process of the near-end algorithm to obtain a differentiable computational flow graph. The prerequisite for implementing this strategy is... The provided fully differentiable algorithm components include linear operators, proximal operators, and proximal algorithms.

[0053] Balanced learning: This paper proposes using balanced learning to improve the high memory usage problem caused by iterative expansion. By treating the iterative process of near-end optimization as a fixed-point iterative problem, both the forward and backward differential processes of multi-step iterations in near-end optimization can be solved by solving a linear equation.

[0054] Unlike iterative expansion methods, this method does not require tracking intermediate variables during the iteration process, thus achieving constant-level memory usage.

[0055] Reinforcement learning: This paper proposes using reinforcement learning to solve the optimization problem of non-differentiable paths. For non-differentiable paths, It is optimized using a model-free reinforcement learning method.

[0056] For example, in optimization problems involving iterative termination conditions in hyperparameter planners, By treating the near-end optimization algorithm as the environment and the hyperparameter planner as the intelligent agent, the interaction between the two can construct the data and optimization objectives required for reinforcement learning, and solve the problem that traditional methods cannot be used for learning because the algorithm is not differentiable.

[0057] Step 4: Perform image reconstruction.

[0058] This invention uses The model constructs a joint optimization algorithm, which includes a differentiable proximal solver and a single thin lens consisting of a parameterized diffraction phase plate.

[0059] Specifically, a Fourier optical imaging model using paraxial optics (i.e., small field of view) and incident plane waves (i.e., objects at optical infinity) is employed. Sensor measurements are represented as the translation-invariant convolution of the image and the point spread function (PSF) parameterized by diffractive optical elements (DOEs), neglecting off-axis aberrations such as coma. Then, the on-axis PSF is derived by simulating the propagation of a wavefront (a plane wave modulated by the wave-correlated phase profile of the DOE) to the sensor plane in free space.

[0060] Furthermore, considering a lens-to-sensor distance of 15 mm and an aperture size of 3 mm, the f-number can be 5. The sensor model has a resolution of 748 × 748 and a pixel pitch of 4 μm. The wavefront propagation simulation is discretized on a 2 μm grid (1496 × 1496) to meet the computational sampling requirements of the transfer function-based Fresnel propagator, while the DOE is modeled as an unconstrained heightmap.

[0061] Specifically, the common design problem of DOE and proximal algorithms for image reconstruction can be represented by the following code:

[0062] x = Variable()

[0063] psf = build_doe_model().psf

[0064] data_term=sum_squares(conv(x,psf)-y)

[0065] reg_term=deep_prior(x,'ffdnet',trainable=True)

[0066] objective=data_term+reg_term

[0067] out=compile(objective,method='admm').solve()

[0068] This invention uses an expanded ADMM solver as... The default training strategy is adopted, and FFDNet (Fast and Flexible Denoising Network) is used as a deep prior to be trained together with DOE, so as to finally realize the image reconstruction of the differentiable proximal algorithm. Example

[0069] The following examples illustrate this. The step-by-step example used. It only requires a few lines of code. This allows various proximal optimization problems to be compiled into various differentiable solvers, such as... Figure 1 As shown, these differentiable solvers can be used independently or integrated into existing differentiable solver pipelines, for example, for optimization with post-processing networks.

[0070] First, consider a joint deconvolution and denoising problem as an example. The forward model for this problem can be defined as follows: The function conv(x,psf) is used to describe this. Given an observation y, the goal is to find the unknown image x by minimizing the objective function. In mathematical expressions, the objective can be concisely described using almost the same syntax as the objective function described mathematically.

[0071] x = Variable()

[0072] data_term=sum_squares(conv(x,psf)-y)

[0073] prior_term=deep_prior(x,unet)

[0074] objective=data_term+prior_term+nonneg(x)

[0075] This facilitates experimentation with different regularization functions. For example, a user can easily change the plug-and-play prior to using the FFDNet denoiser deep_prior(x, 'FFDNet'), or use the total variational prior norm1(grad(x)). The objective can be solved by encapsulating it in a Problem class and calling the solver by specifying the desired algorithm (e.g., ADMM in this case).

[0076] p = Problem(objective)

[0077] out=p.solve(method='admm')

[0078] Allows users to pass more options to configure this process, such as initial solution, maximum number of iterations, and algorithm parameters, p.solve(x0=..., rhos=..., max_iters=...).

[0079] exist In this context, `p.solve` is a differentiable process that can be integrated with PyTorch and used for downstream learning-based tasks. For example... Figure 2 As shown, the compiler-generated solver can propagate gradients to trainable parameters in learning-based components, such as learnable linear operators, proximal regularizers, and internal algorithm parameter schedulers, via backpropagation. Nevertheless, for a training process, using the compiled solver along with placeholders is generally more efficient. To illustrate this, consider a learnable optical diffusion function (PSF), i.e.:

[0080] y = Placeholder(input)

[0081] data_term=sum_squares(conv(x,psf_DOE)-y)

[0082] prior_term=deep_prior(x,unet)

[0083] objective=data_term+prior_term+nonneg(x)

[0084] s=compile(objective,method='admm')

[0085] To jointly train this PSF, a straightforward choice to make the solver differentiable is iterative expansion, which can be achieved directly by fixing the `max_iters` parameter of `p.solve`. For finer-grained control, A primitive `specialize` is provided to specialize an existing solver into a solver for iterative expansion; for example, it can control the use of independent parameters for each iteration.

[0086] s2=specialize(s,method='unroll',shared=False)

[0087] The iterative expansion solver obtained by the above method allows for fine-tuning of the learning-based PnP before or even before training the learning-based linear operator (e.g., the aforementioned learning PSF with the physical imaging model PSF_DOE).

[0088] Iterative unrolling strategies suffer from memory overflow issues for optimization problems requiring a very large number of iterations. Therefore, It supports the integration of deep equilibrium learning across multiple iterations, which from the user's perspective is similar to constructing a differential solver based on iterative expansion using specialized primitives.

[0089] s3=specialize(s,method='deq')

[0090] The DEQ solver offers higher memory efficiency, and its training can be performed by specifying the solver, training dataset, and configuration. The provided helper function performs train_DEQ(s3, dataset, **cfg).

[0091] Differentiable solvers It can also automatically estimate hyperparameters and termination iteration time. Specifically, this method extends the specialized primitives to easily convert the solver into a parameter-tunable solver that can be learned through reinforcement learning.

[0092] s4=specialize(s,method='rl',policy='resnet')

[0093] Training these parameter-free solvers can also be achieved using the built-in training utility train_rl(s4, dataset, **cfg).

[0094] All solvers in the system, including naive and specialized solvers, contain the `s.solve` method, which shares the same function signature. Therefore, users can seamlessly switch between different specialized solvers, providing flexibility for rapid prototyping of various applications.

Claims

1. An image reconstruction method based on a differentiable proximal-end algorithm, characterized in that, First, we model the proximal algorithm for image reconstruction. Then, we design a programming framework for the proximal algorithm for image reconstruction. Next, we perform the differential design of the proximal algorithm for image reconstruction. Finally, we use the proximal algorithm and framework to perform image reconstruction. Step 1: Modeling the near-end algorithm for image reconstruction; The objective of every standard form solution is to find an unknown variable by minimizing an objective function. Let n represent the real number field and n represent the vector dimension. The objective function is defined by the penalty function f. i The sum of these components constitutes the penalty function, with each penalty function taking an unknown variable after a certain linear transformation as its input K. i x, the standard form has a series of additional constraints c j As shown below: Where argmin represents the value of the variable that minimizes the latter expression; st indicates compliance; I represents the number of penalty functions; θ represents the implicit parameter; c represents the constraint condition; and f is the penalty function. J represents any; J represents the number of constraints; K represents any. i For linear operators, Responsible for projecting the unknowns onto the penalty function f i : On the input, constraint c j : Added via equality / inequality, m i This represents the number of rows in the matrix; here, f i y i and c j Through a set of implicit parameters Perform parameterization; Step 2: Design a proximal algorithm programming framework for image reconstruction; Introduction -Prox, a domain-specific modeling language and compiler that addresses memory explosion and manipulability issues in image reconstruction methods using differentiable proximal algorithms; -Prox first defines a placeholder representing the unknown variable x and an observed input y; then, it defines each term in the optimization objective by combining linear operators and proximal functions / operators that may have latent parameters, where linear operators can be arbitrarily combined to form a linear expression, where each leaf node is a variable; each proximal function accepts a linear expression and optional keyword arguments as input according to its specific function signature; after all proximal functions are defined, they are mixed into an optimization objective and wrapped in a problem container to construct the corresponding optimization problem; Step 3: Differential design of the near-end algorithm for image reconstruction; Three differentiation strategies are employed, including iterative expansion, equilibrium learning, and reinforcement learning, among which: Iterative expansion: -Prox supports a differentiation strategy based on iterative expansion; this strategy expands the iterative optimization process of the near-end algorithm to obtain a differentiable computational flow graph; the prerequisite for implementing this strategy is... -Prox provides fully differentiable algorithmic components including linear operators, proximal operators, and proximal algorithms; Balanced learning: -Prox proposes using balanced learning to improve the high memory usage problem caused by iterative expansion; by treating the iterative process of near-end optimization as a fixed-point iterative problem, the forward and backward differential processes of multi-step iterations of near-end optimization can both be solved by solving a linear equation; Reinforcement learning: -Prox proposes using reinforcement learning to solve the optimization problem of non-differentiable paths; for non-differentiable paths, -Prox optimizes it using a model-free reinforcement learning approach; Step 4: Perform image reconstruction; use -The Prox model constructs a jointly optimized proximal algorithm, which includes a differentiable proximal solver and a single thin lens consisting of a parameterized diffraction phase plate. Using an expanded ADMM solver as -Prox's default training strategy, and uses FFDNet as a deep prior trained together with DOE, ultimately achieving image reconstruction using a differentiable proximal algorithm.

2. The image reconstruction method based on a differentiable proximal-end algorithm as described in claim 1, characterized in that, In step 4, the paraxial optics and Fourier optical imaging model of incident plane wave are used to represent the sensor measurement as the translation-invariant convolution result of the image and the point spread function PSF parameterized by the diffractive optical element DOE, and the off-axis aberrations are ignored. Then, the on-axis PSF is derived by simulating the wavefront propagation to the sensor plane in free space.

Images