test类samplePtyLo_Guass(manager,recon_params)初始化之后# 就这一行搞定 # 默认模式去掉白边只有纯原子save_obj_fft(sample.quick_get_object_complex()[0],init_params[workdir]/init_no_bg.png)# 加上 crop_bgFalse保留最原始的白边跟最初的效果一样save_obj_fft(sample.quick_get_object_complex()[0],init_params[workdir]/init_with_bg.png,crop_bgFalse)# 执行重构之前 sample.recon_base(recon_params)# # 下面是你新加的画图代码用来查看 3σ 拼接效果# importmatplotlib.pyplotaspltimportnumpyasnpdefsave_obj_fft(obj_layer,save_path,crop_bgTrue,bg_amp1.0): 极简封装保存 相位振幅FFT 三张图 obj_layer: 一层复数图像 (GPU或CPU格式都行) save_path: 保存的路径如 test.png crop_bg: True去掉白边背景(默认), False保留原始白边背景 bg_amp: 背景振幅值默认 1.0 (仅在 crop_bgTrue 时生效) # 1. 转 CPUifhasattr(obj_layer,get):obj_layerobj_layer.get()phasenp.angle(obj_layer)ampnp.abs(obj_layer)# 2. 根据 flag 决定是否裁剪ifcrop_bg:masknp.abs(amp-bg_amp)0.001rows,colsnp.any(mask,axis1),np.any(mask,axis0)ifnp.any(rows)andnp.any(cols):rnp.where(rows)[0][[0,-1]]cnp.where(cols)[0][[0,-1]]buf3slc(slice(max(0,r[0]-buf),min(amp.shape[0],r[1]buf1)),slice(max(0,c[0]-buf),min(amp.shape[1],c[1]buf1)))phase,ampphase[slc],amp[slc]# 3. 算 FFTfft_lognp.log(np.abs(np.fft.fftshift(np.fft.fft2(amp)))1e-8)# 4. 画图保存fig,axesplt.subplots(1,3,figsize(18,5))axes[0].imshow(phase,cmapjet);axes[0].set_title(Phase)axes[1].imshow(amp,cmapgray);axes[1].set_title(Amplitude)axes[2].imshow(fft_log,cmapinferno);axes[2].set_title(FFT)plt.tight_layout()plt.savefig(save_path,dpi150)plt.close()# 就这一行搞定 # 默认模式去掉白边只有纯原子save_obj_fft(sample.quick_get_object_complex()[0],init_params[workdir]/init_no_bg.png)# 加上 crop_bgFalse保留最原始的白边跟最初的效果一样save_obj_fft(sample.quick_get_object_complex()[0],init_params[workdir]/init_with_bg.png,crop_bgFalse)# class PtyLo(EnginesMixin)iter_000_check.jpgiter_001_check.jpgiter_002_check.jpgdef_engine_iteration_end(self,i,error,start_time):self.update_manager()self.manager._engine_iteration_end(i,error,start_time)# # 新增在指定迭代次数时保存中间过程图# # 设定你想看哪几次的结果比如第 0 次初始、第 1 次、第 5 次、第 10 次watch_iterations[0,1,2,3,4,5,10,20,50]ifiinwatch_iterations:importmatplotlib.pyplotaspltimportnumpyasnp# 1. 获取当前的复数物体 (记得加 self._get_object_complex() 或者直接用 self.manager._object)# 注意这里必须用 manager 里的 object因为 update_manager 刚刚同步过最新状态current_objself.manager._object# 2. 取第 0 层转 CPUlayer_0current_obj[0]ifhasattr(layer_0,get):layer_0layer_0.get()phase_imgnp.angle(layer_0)amp_imgnp.abs(layer_0)# 3. 画图fig,axesplt.subplots(1,2,figsize(12,5))im1axes[0].imshow(phase_img,cmapjet)axes[0].set_title(fIter{i}- 相位图 (Error:{error:.4f}))plt.colorbar(im1,axaxes[0])im2axes[1].imshow(amp_img,cmapgray)axes[1].set_title(fIter{i}- 振幅图)plt.colorbar(im2,axaxes[1])plt.tight_layout()# 4. 保存到工作目录不弹窗直接存盘save_pathself.manager.workdir/fiter_{i:03d}_check.pngplt.savefig(save_path,dpi150)plt.close()# 关闭画板释放显存极其重要否则多张图会把显存撑爆#
