From 1dd389d08a159807c77074d9d0df98720ece8c80 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 28 May 2025 10:09:02 +0200 Subject: [PATCH] Update the initialization for near field to work better with oversampling --- .../tools/image_processing/image_processing.py | 16 ++++++++++------ src/cdtools/tools/initializers/initializers.py | 13 ++++--------- 2 files changed, 14 insertions(+), 15 deletions(-) diff --git a/src/cdtools/tools/image_processing/image_processing.py b/src/cdtools/tools/image_processing/image_processing.py index 8a9eb4b..b774daa 100644 --- a/src/cdtools/tools/image_processing/image_processing.py +++ b/src/cdtools/tools/image_processing/image_processing.py @@ -321,19 +321,23 @@ def convolve_1d(image, kernel, dim=0, fftshift_kernel=True): return conv_im -def fourier_upsample(ims, preserve_mean=False): +def fourier_upsample(ims, upsample_factor=2, preserve_mean=False): # If preserve_mean is true, it preserves the mean pixel intensity # otherwise, it preserves the total summed intensity - upsampled = t.zeros(ims.shape[:-2]+(2*ims.shape[-2],2*ims.shape[-1]), + + upsampled = t.zeros(ims.shape[:-2]+(upsample_factor*ims.shape[-2], + upsample_factor*ims.shape[-1]), dtype=ims.dtype, device=ims.device) - left = [ims.shape[-2]//2,ims.shape[-1]//2] - right = [ims.shape[-2]//2+ims.shape[-2], - ims.shape[-1]//2+ims.shape[-1]] + + left = [((upsample_factor-1)*ims.shape[-2])//2, + ((upsample_factor-1)*ims.shape[-1])//2] + right = [left[0]+ims.shape[-2], + left[1]+ims.shape[-1]] upsampled[...,left[0]:right[0],left[1]:right[1]] = propagators.far_field(ims) if preserve_mean: - upsampled *= 2 + upsampled *= upsample_factor return propagators.inverse_far_field(upsampled) diff --git a/src/cdtools/tools/initializers/initializers.py b/src/cdtools/tools/initializers/initializers.py index 1925b95..c3aa47e 100644 --- a/src/cdtools/tools/initializers/initializers.py +++ b/src/cdtools/tools/initializers/initializers.py @@ -419,15 +419,10 @@ def SHARP_style_near_field_probe(dataset, backward_propagator, oversampling=1): probe_guess = backward_propagator(probe_guess_det_plane) if oversampling != 1: - probe_guess_real = t.nn.functional.interpolate( - t.real(probe_guess).unsqueeze(0).unsqueeze(0), - scale_factor=oversampling, mode='bilinear')[0,0] - probe_guess_imag = t.nn.functional.interpolate( - t.imag(probe_guess).unsqueeze(0).unsqueeze(0), - scale_factor=oversampling, mode='bilinear')[0,0] - - probe_guess = ( t.complex(probe_guess_real, probe_guess_imag) - / oversampling**2 ) + probe_guess = image_processing.fourier_upsample( + probe_guess, + upsample_factor=oversampling, preserve_mean=False + ) return probe_guess