EfficienTIE

TIE（Transport of Intensity Equation）的神经网络求解器，物理层源自论文On a universal solution to the transport-of-intensity equation。

功能特性

• TIE快速求解
• 现场数据生成
• 不同相位范围估计
• 帧数评测

文件结构

.
├─ best_model.pth      # 预训练权重
├─ efficientie.py      # 神经网络
├─ make_data.py        # 数据生成器
├─ speed.py            # 帧数测试
├─ test.py             # 测试集测试
└─ train.py            # 训练代码

核心类

EfficenTIE

class EfficenTIE(nn.Module):
    def __init__(self, in_channels=3, base_c=32):
        super().__init__()
        
        self.physics = Physics()
        
        self.stem = ConvBNReLU(in_channels, base_c, kernel_size=3, padding=1)
        
        self.enc1 = ResBlock(base_c)
        self.down1 = ConvBNReLU(base_c, base_c*2, stride=2)
        
        self.enc2 = ResBlock(base_c*2)
        
        self.global_corrector = GlobalContext(base_c*2, reduction=4)
        
        self.up1 = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear'),
            ConvBNReLU(base_c*2, base_c, 1)
        )
        
        self.dec1 = ResBlock(base_c)
        self.fusion = ConvBNReLU(base_c * 2, base_c, 1)
        
        self.tail = nn.Sequential(
            nn.Conv2d(base_c, base_c, 3, padding=1, bias=False),
            nn.BatchNorm2d(base_c),
            nn.ReLU(inplace=True),
            nn.Conv2d(base_c, 1, 3, padding=1),
            nn.Tanh()
        )

        
    def forward(self, I, dI_dz, wavelength, dx, dy):
        delta_dIdz = dI_dz.clone()
        
        Phi = self.physics(I, delta_dIdz, wavelength, dx, dy, device=I.device)
    
        x_in = torch.cat([Phi, I, dI_dz], dim=1)

        # 编码器
        x0 = self.stem(x_in)      # [B, 32, H, W]
        x1 = self.enc1(x0)        # [B, 32, H, W]
        x1_down = self.down1(x1)  # [B, 64, H/2, W/2]
        
        x2 = self.enc2(x1_down)   # [B, 64, H/2, W/2]
        
        # 瓶颈层
        x2_corrected = self.global_corrector(x2)
        
        # 解码器
        u1 = self.up1(x2_corrected) # [B, 32, H, W]
        
        # 跳跃连接融合
        concat = torch.cat([u1, x1], dim=1) # [B, 64, H, W]
        fused = self.fusion(concat)         # [B, 32, H, W]
        out_feat = self.dec1(fused)
        
        # 输出
        delta_phi = self.tail(out_feat)

        phi = Phi + delta_phi
        phi = phi - phi.mean(dim=(-2, -1), keepdim=True)

        return phi

快速开始

使用预训练模型

import torch
from efficentie import EfficenTIE

model = EfficenTIE(in_channels=3, base_c=32)
checkpoint = torch.load('best_model.pth', map_location='cuda')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval().cuda()

I = torch.randn(1, 1, 256, 256).cuda()  #[B, 1, H, W]
dI_dz = torch.randn(1, 1, 256, 256).cuda()  #[B, 1, H, W]
lam = torch.tensor(620e-9).cuda()  #(m)
dx = dy = torch.tensor(1.85e-6).cuda()  #(m)

with torch.no_grad():
    phi = model(I, dI_dz, lam, dx, dy)  # [1, 1, 256, 256]

完整测试示例

python test.py

输入输出

输入

• 光强图和其对应的光强导数
• 波长、像素尺寸、等物理参数

输出

• 恢复的相位图

依赖

• PyTorch
• NumPy
• matplotlib
• torchmetrics

引用

如果您在研究中使用了这段代码（当然，这不太可能），请引用原论文：

• BibTeX:

  @article{Zhang:20,
    author={Jialin Zhang and Qian Chen and Jiasong Sun and Long Tian and Chao Zuo},
    journal = {Opt. Lett.},
    keywords = {Fourier transforms; Imaging systems; Microlens arrays; Optical fields; Phase imaging; Phase retrieval},
    number = {13},
    pages = {3649--3652},
    publisher = {Optica Publishing Group},
    title = {On a universal solution to the transport-of-intensity equation},
    volume = {45},
    month = {Jul},
    year = {2020},
    url = {https://opg.optica.org/ol/abstract.cfm?URI=ol-45-13-3649},
    doi = {10.1364/OL.391823}
  

}

碎碎念

• 孩子们，逆问题太难了。测试集SSIM才0.87 ( ⩌ - ⩌ )
• 这个网络也可以叫GBC-TIE(Global Bias Correction)