论文复现与解析(neural holography)

项目结构

  • main.py:根据不同的相位优化算法(SGD/GS/DPAC/Holonet/U-net)和模型(ASM/CITL)来生成相位图

  • eval.py:根据模型(ASM/CITL)来评估相位图和重建图像

  • main_eval.sh:一个bash脚本,带一个传入参数。首先在RGB通道上根据不同的相位优化算法(SGD/GS/DPAC/Holonet/U-net)执行main.py,然后在RGB通道上根据不同的相位优化算法(SGD/GS/DPAC/Holonet/U-net)执行eval.py

  • propagation_ASM.py:实现角谱法相关算子

  • propagation_model.py:定义CITL相关模型

  • holenet.py:定义HoloNet/U-Net相关模型

  • algorithms.py:实现相位提取算法(GS/SGD/DPAC)

  • train_holonet.py:训练HoloNet

  • train_model.py:训练CITL模型

  • ./data/:测试数据保存位置

  • ./pretrained_networks/:HoloNet/UNet在RGB通道的预训练模型保存位置

./utils/:工具函数实现

  • utils.pycontains utility functions.

  • modules.pycontatins PyTorch wrapper modules for easy use of algorithms.pyand our hardware controller.

  • pytorch_prototyping/submodule contains custom pytorch modules with sane default parameters. (adapted from here)

  • augmented_image_loader.pycontains modules of loading a set of images.

  • utils_tensorboard.py contains utility functions used for visualization on tensorboard.

  • slm_display_module.py contains the SLM display controller module. (HOLOEYE SDK)

  • detect_heds_module_path.pysets the SLM SDK path. Otherwise you can copy the holoeye module directory into your project and import by using import holoeye.

  • camera_capture_module.py contains the FLIR camera capture controller module. (PyCapture2 SDK)

  • calibration_module.py contains the homography calibration module.

前期准备

git clone [email protected]:computational-imaging/neural-holography.git
  • 创建虚拟环境(按照environment_windows.yml安装,名字为neural-holography,并安装需要的包)
conda env create -f environment_windows.yml
conda activate neural-holography

  • 问题汇总
 问题:pycapture2安装不了
# 描述:pycapture2是相机拍照的包,可以手动安装
# 解决:进入flycapture skd官网(https://www.flir.cn/products/flycapture-sdk/),
# china->现在下载->...->FlyCapture2->Windows->Python->PyCapture2-2.13.61.win-amd64-py3.6.msi改成(PyCapture2-2.13.61.win32-py3.6)下载

问题:holoeye python sdk没有安装
# 描述:官网上的sdk需要设备、账号才能下载,而且下载不下来
# 解决:用1.0版本的python sdk安装(代理商给的),并且改一下环境变量的名字(HEDS_2_PYTHON_MODULES、HEDS_2_PYTHON),其中的2代表版本,目前一共有三个版本,这里论文中使用的为2,而安装的为1,所以需要一下。

问题:某些modules没有、找不到
# 描述:不知道为什么
# 解决:使用pip install **安装

问题:pytorch_prototyping下的文件没有、找不到
# 描述:该文件夹下的代码不是本论文作者写的
# 解决:去github上找到该文件单独拷贝下来,然后粘贴到相应的文件夹下

问题:pretrained_networks下模型打不开
# 描述:需要解压
# 解决:解压文件夹,将解压后的文件夹内容放到pretrained_networks下

问题:cuda存储空间不够
# 描述:代码使用GPU加速,显存不够或者batchsize太大
# 解决:增大显存或者改小batchsize

问题:import torch 报错[winerror 182]load error:......\caffe2_detectron_ops_gpu.dll...
# 描述:不知道为什么,可能是python、torch、cudatoolkit或者其他包之间的版本兼容性问题,或者是多次import torch(本文件、其他文件、import 其他文件(应该不是这个:我直接import torch也会报错))
# 可能解决:1. 删掉这些dll文件
# 2. pip install intel-openmp
# 3. 尝试降低torch及依赖的版本到torch==1.4.0,但是存在问题: fft模块在torch新旧版本不同,torch1.4.0不支持complex操作,
# 4. 尝试将所有包安装默认版本(删掉environment_windows.yml中指定版本的部分)/安装python3.7/3.8,存在问题:错误依旧、某些第三方包(pycapture2)只支持到python3.6
# 5. 尝试观察代码内部的问题,注意到:在python命令行交互模式下import torch和训练holonet时会报上述错误,但是一开始做相位优化和重建时,同样import torch却没有报错,这是什么意思呢?尝试在train_holonet.py下观察代码,发现文件内调用了torch但是import torch是灰色的。从这点出发可能找到问题所在
# 解决?:删掉灰色部分,重新运行代码,正常跑起来了(环境仍然和environment_windows.yml中一致)

更新:2022.2.26 import torch报错的问题
# 1. 怀疑是FlyPyCapture2包和其他包之前不相互兼容或者依赖存在问题:我按照environment_windows.yml文件创建环境但是不安装PyCapture2,此时import torch是正常的,但是安装了PyCapture2之后出现上述错误
# 2. 暂时不使用相机的相关功能,代码运行正常(不安装PyCapture2)
# 3. 使用pytorch==1.4,作者在readme中讲使用utils/tuils.py的复数功能,但是不知道如何做
# 4. 或许等到相机来了之后,情况会不一样?
# 5. 使用Ubuntu?还没去做试验

更新:2022.3.1 import torch报错的问题
# 1. 之前下载的amd64和win32版本有问题,当时认为amd64是给64位操作系统,win32是32位操作系统
# 2. 解决方案:卸载amd64版本,安装win32版本,能正常import torch
# 3. 不知道为什么

更新:2022.7.13 关于holoeye slm sdk
# 在官网上下载了windows python sdk 3.0版本,需要修改部分内容:
# 1. 从example目录下找到detect_heds_module_path.py文件,因为为了避免在每个python环境下安装sdk所以使用了环境变量的形式,需要使用这个脚本来避免import slm有关API报错。此外,如果现在直接运行,可能会报错,仔细一看detect_heds_module_path.py 发现,里面默认环境变量的命名和我们sdk安装的默认命名不一样,修改成一样即可(去掉_3)。
# 2. 修改SLMDisplay下的方法,因为sdk版本不同,推荐使用的API也不同

演示

在这里插入图片描述

Wave Propagation Model

角谱法(ASM)

  • 角谱法通过传递函数来实现,即变换到频域,然后乘以传递函数,随后变换到空域
  • 角谱法传播波场关键代码:
def propagation_ASM(u_in, feature_size, wavelength, z, linear_conv=True,
                    padtype='zero', return_H=False, precomped_H=None,
                    return_H_exp=False, precomped_H_exp=None,
                    dtype=torch.float32):
    """Propagates the input field using the angular spectrum method

    Inputs
    ------
    u_in: PyTorch Complex tensor (torch.cfloat) of size (num_images, 1, height, width) -- updated with PyTorch 1.7.0
    feature_size: (height, width) of individual holographic features in m
    wavelength: wavelength in m
    z: propagation distance
    linear_conv: if True, pad the input to obtain a linear convolution
    padtype: 'zero' to pad with zeros, 'median' to pad with median of u_in's
        amplitude
    return_H[_exp]: used for precomputing H or H_exp, ends the computation early
        and returns the desired variable
    precomped_H[_exp]: the precomputed value for H or H_exp
    dtype: torch dtype for computation at different precision

    Output
    ------
    tensor of size (num_images, 1, height, width, 2)
    """

    if linear_conv:
        # preprocess with padding for linear conv.
        input_resolution = u_in.size()[-2:]
        conv_size = [i * 2 for i in input_resolution]
        if padtype == 'zero':
            padval = 0
        elif padtype == 'median':
            padval = torch.median(torch.pow((u_in**2).sum(-1), 0.5))
        u_in = utils.pad_image(u_in, conv_size, padval=padval, stacked_complex=False)

    if precomped_H is None and precomped_H_exp is None:
        # resolution of input field, should be: (num_images, num_channels, height, width, 2)
        field_resolution = u_in.size()

        # number of pixels
        num_y, num_x = field_resolution[2], field_resolution[3]

        # sampling inteval size
        dy, dx = feature_size

        # size of the field
        y, x = (dy * float(num_y), dx * float(num_x))

        # frequency coordinates sampling
        fy = np.linspace(-1 / (2 * dy) + 0.5 / (2 * y), 1 / (2 * dy) - 0.5 / (2 * y), num_y)
        fx = np.linspace(-1 / (2 * dx) + 0.5 / (2 * x), 1 / (2 * dx) - 0.5 / (2 * x), num_x)

        # momentum/reciprocal space
        FX, FY = np.meshgrid(fx, fy)

        # transfer function in numpy (omit distance)
        HH = 2 * math.pi * np.sqrt(1 / wavelength**2 - (FX**2 + FY**2))

        # create tensor & upload to device (GPU)
        H_exp = torch.tensor(HH, dtype=dtype).to(u_in.device)

        ###
        # here one may iterate over multiple distances, once H_exp is uploaded on GPU

        # reshape tensor and multiply
        H_exp = torch.reshape(H_exp, (1, 1, *H_exp.size()))

    # handle loading the precomputed H_exp value, or saving it for later runs
    elif precomped_H_exp is not None:
        H_exp = precomped_H_exp

    if precomped_H is None:
        # multiply by distance
        H_exp = torch.mul(H_exp, z)

        # band-limited ASM - Matsushima et al. (2009)
        fy_max = 1 / np.sqrt((2 * z * (1 / y))**2 + 1) / wavelength
        fx_max = 1 / np.sqrt((2 * z * (1 / x))**2 + 1) / wavelength
        H_filter = torch.tensor(((np.abs(FX) < fx_max) & (np.abs(FY) < fy_max)).astype(np.uint8), dtype=dtype)

        # get real/img components
        H_real, H_imag = utils.polar_to_rect(H_filter.to(u_in.device), H_exp)

        H = torch.stack((H_real, H_imag), 4)
        H = utils.ifftshift(H)
        H = torch.view_as_complex(H)
    else:
        H = precomped_H

    # return for use later as precomputed inputs
    if return_H_exp:
        return H_exp
    if return_H:
        return H

    # For who cannot use Pytorch 1.7.0 and its Complex tensors support:
    # # angular spectrum
    # U1 = torch.fft(utils.ifftshift(u_in), 2, True)
    #
    # # convolution of the system
    # U2 = utils.mul_complex(H, U1)
    #
    # # Fourier transform of the convolution to the observation plane
    # u_out = utils.fftshift(torch.ifft(U2, 2, True))

    U1 = torch.fft.fftn(utils.ifftshift(u_in), dim=(-2, -1), norm='ortho')

    U2 = H * U1

    u_out = utils.fftshift(torch.fft.ifftn(U2, dim=(-2, -1), norm='ortho'))

    if linear_conv:
        # return utils.crop_image(u_out, input_resolution) # using stacked version
        return utils.crop_image(u_out, input_resolution, pytorch=True, stacked_complex=False)  # using complex tensor
    else:
        return u_out

相机在环校准模型(CITL)

数据集

模型的训练

测试

Phase Optimization Method

GS算法

  • 测试:python main.py --channel=0 --method=GS --root_path=./phases

SGD算法

  • 测试:
    • SGD with ASM(Gradient Descent):python main.py --channel=0 --method=SGD --root_path=./phases
    • SGD with CITL(Gradient Descent):python main.py --channel=0 --method=SGD --prop_model='MODEL' --prop_model_dir=YOUR_MODEL_PATH --root_path=./phases

CITL优化算法

  • 测试:python main.py --channel=0 --method=SGD --citl=True --root_path=./phases

DPAC算法

  • 测试:python main.py --channel=0 --method=DPAC --root_path=./phases

U-Net优化算法

  • 测试:python main.py --channel=0 --method=UNET --root_path=./phases --generator_dir=./pretrained_networks

HoloNet模型

数据集

  • Phase pattern(Generated by iterative algorithm) + Target intensity image

模型的训练

测试

  • Training with ASM:python train_holonet.py --perfect_prop_model=True --run_id=my_first_holonet --batch_size=4 --channel=0

猜你喜欢

转载自blog.csdn.net/BubbleCodes/article/details/122920344
今日推荐