[HELP] Image alignment for SSIM calculation

[fu0601]

Hi,

I noticed that after several updates, ptyrad has become much faster. However, in my recent tests, I found that I could not reproduce the results shown in Figure 3a of your paper, as shown in the attached figure.

I have also attached the script parameters I used for the run.

ptyrad version: 0.1.0b13
nvcc: CUDA 12.4
GPU: NVIDIA RTX 4090 24G

[sim_tBL_WSe2.yml](https://github.com/user-attachments/files/26052174/sim_tBL_WSe2.yml)

[chiahao3]

Hi @fu0601 ,

Thanks for trying out PtyRAD! Yes, PtyRAD has become significantly faster since last summer with the JIT compilation if you set compiler_configs: {'enable': true} under recon_params, and I did a recent benchmark here showing a 10x speedup with compiled PtyRAD vs. PtyShv under the same condition. See doc for example params of how to turn on JIT compilation.

If you get a chance, maybe you can upload the depth-summed reconstructed phase image here? I don't think it's a convergence issue as I've run the demo reconstructions (tBL-WSe2 and PSO) multiple times before every release to make sure the reconstruction quality is unaffected.

I think the drop in SSIM is almost certainly an image alignment issue, as the plotting notebook contains a hard-coded image shift value (shift = [-0.205,-0.2675]) obtained manually via imageJ. This shift is necessary as position correction can introduced a global shift between the ground truth and the reconstructed image, and these shifts are different for each package and needs to be determined for every run unless you set the random seed. If you get a chance to fine-tune the shift vector a little bit, I think you'll find the SSIM swinging between 0.8 to 0.9.

Executed cell [6]

crop_height, crop_width = 170, 170 # Note that the window size and alignment would also impact the SSIM value
# The image shifts are obtained via SIFT registration using imageJ

# Get images
gt_phase = center_crop(load_hdf5(gt_path, key = 'gt_phase').sum(0), crop_height, crop_width, offset = [1,1])

ssims_ptyrad_c = []
for img_ptyrad_c in objs_ptyrad_c:
    img_ptyrad_c = center_crop(img_ptyrad_c, crop_height, crop_width, offset = [0,0])
    shift = [-0.205,-0.2675] # Check this <=========== This will change for every run unless you fix the random seed
    img_ptyrad_c = np.real(np.fft.ifftn(fourier_shift(np.fft.fftn(img_ptyrad_c), shift)))
    img_ptyrad_c -= img_ptyrad_c.min()
    
    ssim_ptyrad_c = ssim(img_ptyrad_c, gt_phase, data_range=(gt_phase.max() - gt_phase.min()))
    print(ssim_ptyrad_c)
    ssims_ptyrad_c.append(ssim_ptyrad_c)

[fu0601]

Hi,
Thank you for your suggestion. My test revealed that it was indeed an issue with image alignment. I am not familiar with the use of ImageJ software. If possible, could you provide a tutorial on how to calculate the image shift value using ImageJ? I would be very grateful. I used the phase_cross_correlation function in skimage.registration library to calculate the image shift value. For the results of the ptyrad experiment, the shift calculated by phase_cross_correlation is not accurate, but for PtyShv, an accurate shift value can be obtained.
I have placed my test results at https://github.com/fu0601/ptyrad_vs_PtyShv.git
`from skimage.registration import phase_cross_correlation
shift, error, _ = phase_cross_correlation(gt_phase, center_crop(objs_ptyrad[-1], crop_height, crop_width, offset = [0,0]), upsample_factor=100)
ssims_ptyrad = []
for img_ptyrad in objs_ptyrad:
img_ptyrad = center_crop(img_ptyrad, crop_height, crop_width, offset = [0,0])
# shift = [-0.205,-0.2675] # Check this
img_ptyrad = np.real(np.fft.ifftn(fourier_shift(np.fft.fftn(img_ptyrad), shift)))
img_ptyrad -= img_ptyrad.min()
ssim_ptyrad = ssim(img_ptyrad, gt_phase, data_range=(gt_phase.max() - gt_phase.min()))
ssims_ptyrad.append(ssim_ptyrad)

shift, error, _ = phase_cross_correlation(gt_phase, center_crop(objs_ptyshv[-1], crop_height, crop_width, offset = [-1,-1]), upsample_factor=100)
ssims_ptyshv = []
for img_ptyshv in objs_ptyshv:
img_ptyshv = center_crop(img_ptyshv, crop_height, crop_width, offset = [-1,-1])
# shift = [-0.35,-0.416]
img_ptyshv = np.real(np.fft.ifftn(fourier_shift(np.fft.fftn(img_ptyshv), shift)))
img_ptyshv -= img_ptyshv.min()
ssim_ptyshv = ssim(img_ptyshv, gt_phase, data_range=(gt_phase.max() - gt_phase.min()))
ssims_ptyshv.append(ssim_ptyshv)`

[chiahao3]

Hi @fu0601 ,

Glad to see that it was just the image alignment!

I won't do an imageJ tutorial since it was just a one-time fix with handy interface, not necessarily a proper solution worth sharing. Ideally we do all required preprocessing in python for reproducibility. You're free to experiment with whatever image alignment approaches, most LLMs should be able to give reasonable solutions. I'll just leave some notes and let you play with it.

1. Image alignment and background level adjustments are both necessary for meaningful SSIM calculation
2. Since position correction can introduce some image distortion, you'll want to crop the image first so your alignment is not affected by the edges, as edges are less constrained by data and generally suffer more distortion
3. Shifting images with Fourier methods can introduce some edge artifacts with unexpected values. For example, you may get negative values at the edge when you shift PtyRADc that contains only positive value. This would affect the later img -= img.min(). I didn't really notice this until last Nov, but a better approach is to crop out the edges again after you shift the image.
4. Phase images can have arbitrary background level, but for SSIM calculation, you'll need to align each image with the ground truth. There's no concensus on how to best align it yet, but I think we can simply align the minimal value AFTER you handle the edges properly. Note that you might want to shift the background level for ground truth as well.

I'd suggest look closer into each step. Check your shifts value, and check the shifted images (both their alignment and value range). It's possible that the phase_cross_correlation found a totally valid shift, but the edges value somehow messed up the base level alignment and affect the SSIM.