PSF: created class GtPSF, for PSF applications (internally uses OTFs)

Signed-off-by: Nicola Vigano <nicola.vigano@esrf.fr>

5_reconstruction/gtAstra3DTV.m

@@ -74,8 +74,16 @@ function [volume, res_norm] = gtAstra3DTV(proj, varargin)
 
     astra_projector_id = astra_create_projector('cuda3d', proj_geom, vol_geom);
 
-    fp = @(x)astra_mex_direct_c('FP3D', astra_projector_id, x);
-    bp = @(x)astra_mex_direct_c('BP3D', astra_projector_id, x);
+    if (isfield(proj, 'psf') && ~isempty(proj.psf))
+        psf_obj = GtPSF();
+        psf_obj.set_psf_direct(proj.psf, [proj.num_cols, proj.num_rows]);
+
+        fp = @(x)psf_obj.apply_psf_direct(astra_mex_direct_c('FP3D', astra_projector_id, x));
+        bp = @(x)astra_mex_direct_c('BP3D', astra_projector_id, psf_obj.apply_psf_transpose(x));
+    else
+        fp = @(x)astra_mex_direct_c('FP3D', astra_projector_id, x);
+        bp = @(x)astra_mex_direct_c('BP3D', astra_projector_id, x);
+    end
 
     bp_weights = bp(ones(size(proj.stack), 'single'));
     fp_weights = fp(ones(size(bp_weights), 'single'));

zUtil_Imaging/GtPSF.m

+classdef GtPSF < handle
+    properties
+        psf_direct = [];
+        psf_transpose = [];
+
+        otf_direct = [];
+        otf_transpose = [];
+
+        is_symmetric = true;
+        is_projection_invariant = true;
+        image_size = [];
+
+        data_type = 'single';
+    end
+
+    methods (Access = public)
+        function set_psf_direct(self, psf_d, image_size)
+            self.reset();
+
+            % let's be sure that it will be in ASTRA's projection
+            % convention
+            if (size(psf_d, 1) == size(psf_d, 2))
+                psf_d = permute(psf_d, [1 3 2]);
+            elseif (size(psf_d, 1) ~= size(psf_d, 3))
+                error([mfilename ':wrong_argument'], ...
+                    'PSFs should be in the form of _square images_, with odd edges')
+            elseif (mod(size(psf_d, 1), 2) ~= 1)
+                error([mfilename ':wrong_argument'], ...
+                    'PSFs should be in the form of square images, with _odd edges_')
+            end
+
+            % let's renormalize the PSFs
+            psfs_norm = 1 ./ sum(sum(abs(psf_d), 1), 3);
+            psf_d = bsxfun(@times, psf_d, psfs_norm);
+
+            self.psf_direct = psf_d;
+            if (exist('image_size', 'var') && ~isempty(image_size))
+                self.image_size = image_size;
+            end
+
+            % Let's find out whether they are symmetric or not
+            psf_t = permute(psf_d, [3 2 1]);
+            self.is_symmetric = all(abs(psf_d(:) - psf_t(:)) < eps('single'));
+            if (~self.is_symmetric)
+                self.psf_transpose = psf_t;
+            end
+
+            self.is_projection_invariant = size(psf_d, 2) == 1;
+
+            % if we know the images sizes, we can already compute the OTF
+            if (~isempty(image_size))
+                self.init_otfs();
+            end
+        end
+
+        function imgs = apply_psf_direct(self, imgs)
+            self.check_incoming_images(imgs);
+
+            imgs = self.apply_otf(imgs, true);
+        end
+
+        function imgs = apply_psf_transpose(self, imgs)
+            self.check_incoming_images(imgs);
+
+            imgs = self.apply_otf(imgs, false);
+        end
+    end
+
+    methods (Access = private)
+        function reset(self)
+            self.psf_direct = [];
+            self.psf_transpose = [];
+
+            self.otf_direct = [];
+            self.otf_transpose = [];
+
+            self.image_size = [];
+        end
+
+        function otf = init_single_otf(self, psf)
+            num_psfs = size(psf, 2);
+            psf_edge = (size(psf, 1) - 1) / 2;
+
+            otf = zeros([self.image_size(1), num_psfs, self.image_size(2)], self.data_type);
+
+            center_otf = floor(self.image_size / 2) + 1;
+            lims_otf_lower = center_otf - psf_edge;
+            lims_otf_upper = center_otf + psf_edge;
+
+            otf(lims_otf_lower(1):lims_otf_upper(1), :, lims_otf_lower(2):lims_otf_upper(2)) = psf;
+
+            otf = ifftshift(otf, 1);
+            otf = ifftshift(otf, 3);
+
+            otf = fft(otf, [], 1);
+            otf = fft(otf, [], 3);
+
+            otf = fftshift(otf, 1);
+            otf = fftshift(otf, 3);
+        end
+
+        function init_otfs(self)
+            self.otf_direct = self.init_single_otf(self.psf_direct);
+
+            if (~self.is_symmetric)
+                self.otf_transpose = self.init_single_otf(self.psf_transpose);
+            end
+        end
+
+        function check_incoming_images(self, imgs)
+            num_imgs = size(imgs, 2);
+            if (~self.is_projection_invariant && (num_imgs ~= size(self.psf_d, 2)))
+                error([mfilename ':wrong_argument'], ...
+                    'The number of images (%d) differs from the number of projection dependent PSFs (%d)', ...
+                    num_imgs, size(self.psf_d, 2))
+            end
+
+            imgs_size = [size(imgs, 1), size(imgs, 3)];
+            if (isempty(self.image_size))
+                self.otf_direct = [];
+                self.otf_transpose = [];
+
+                self.image_size = imgs_size;
+            elseif (any(self.image_size ~= imgs_size))
+                warning([mfilename ':wrong_argument'], ...
+                    'OTFs computed for the wrong image size ([%d, %d] instead of [%d, %d]). Recomputing them...', ...
+                    self.image_size, imgs_size)
+
+                self.otf_direct = [];
+                self.otf_transpose = [];
+
+                self.image_size = imgs_size;
+            end
+
+            if (isempty(self.otf_direct))
+                self.init_otfs();
+            end
+        end
+
+        function imgs = apply_otf(self, imgs, is_direct)
+            imgs = ifftshift(imgs, 1);
+            imgs = ifftshift(imgs, 3);
+
+            imgs = fft(imgs, [], 1);
+            imgs = fft(imgs, [], 3);
+
+            imgs = fftshift(imgs, 1);
+            imgs = fftshift(imgs, 3);
+
+            if (is_direct || self.is_symmetric)
+                imgs = bsxfun(@times, imgs, self.otf_direct);
+            else
+                imgs = bsxfun(@times, imgs, self.otf_transpose);
+            end
+
+            imgs = ifftshift(imgs, 1);
+            imgs = ifftshift(imgs, 3);
+
+            imgs = ifft(imgs, [], 1);
+            imgs = ifft(imgs, [], 3);
+
+            imgs = fftshift(imgs, 1);
+            imgs = fftshift(imgs, 3);
+
+            imgs = real(imgs);
+        end
+    end
+end
\ No newline at end of file