光学仪器  2024, Vol. 46 Issue (5): 51-57   PDF    
基于Trans-MIR模型的多模态图像重建
李一鸣1, 王昊2, 李然2, 陈泉2, 卢海军3, 杨晖1,2     
1. 上海健康医学院 医疗器械学院,上海 201318;
2. 上海理工大学 光电信息与计算机工程学院,上海 200093;
3. 上海诺基亚贝尔股份有限公司,上海 201206
摘要: 图像重建是光学计算成像的关键环节之一。目前基于深度学习的图像重建主要使用卷积神经网络、循环神经网络或生成对抗网络等模型。大多数研究仅通过单一模态的数据训练模型,难以在保证成像质量的同时又具备不同场景的泛化能力。为解决这一问题,提出了一种基于Transformer模块的多模态图像重建模型(multi-modal image reconstruction model based on the Transformer, Trans-MIR)。实验结果表明,Trans-MIR能够从多模态数据中提取图像特征,实现高质量的图像重建,对二维通用人脸散斑图像进行图像重建的结构相似度高达0.93,对三维微管结构图像的超分辨重建的均方误差低至10−4量级。Trans-MIR对研究多模态图像重建具有一定的启发作用。
关键词: 图像重建    光学计算成像    多模态    Transformer模块    
Multi-modal image reconstruction method based on Trans-MIR model
LI Yiming1, WANG Hao2, LI Ran2, CHEN Quan2, LU Haijun3, YANG Hui1,2     
1. School of Medical Instruments, Shanghai University of Medicine and Health Sciences, Shanghai 201318, China;
2. School of Optical-Electrical and Computer Engineering, University of Shanghai for Science and Technology, Shanghai 200093, China;
3. Nokia Shanghai Bell Co., LTD., Shanghai 201206, China
Abstract: Image reconstruction is one of the key steps in the optical computational imaging. At present, image reconstruction based on deep learning mainly uses convolutional neural network, cyclic neural network and generative adversarial network. Most models are only trained through the data of a single mode, which is difficult to ensure the quality of imaging while possessing the generalization ability of different scenes. To solve this problem, a multi-modal image reconstruction model based on the Transformer (Trans-MIR) is proposed in this paper. Experimental results show that Trans-MIR can extract image features from multi-modal data to achieve high-quality image reconstruction. The structural similarity of 2D universal face speckle reconstruction was as high as 0.93 and the mean square error of 3D microtubule reconstruction was as low as 10−4. It provides inspiration for the study of multimodal image reconstruction.
Key words: image reconstruction    optical computational imaging    multi-modal    Transformer module    
引 言

随着深度学习与计算机视觉技术的不断发展,光学计算成像技术受到了越来越多的关注。光学计算成像使用计算机技术对成像过程中的光学系统、传输介质、探测器等关键部分进行建模和优化,突破了传统光学成像技术的限制。计算成像具有分辨率高,探测距离远,视场大,成本低和灵活性好等特点,已经广泛应用在医疗影像、工业检测、航空航天、智能交通等领域[1-2]。图像重建是光学计算成像的核心应用,能够从接收到的失真信号中恢复出原始目标物的图像。对图像重建技术的研究已经成为当前计算成像重要的一环,对医疗民生和国家战略等方面具有重要意义。

当前,图像重建已经形成了一些比较成熟的技术路线和算法,如压缩感知、小波变换、全变分正则化、支持向量回归和深度学习等方法。基于深度学习的图像重建方法主要使用卷积神经网络(convolutional neural network,CNN)、循环神经网络(recurrent neural network,RNN)和生成对抗网络(generative adversarial network,GAN)等技术,可以直接从数据中学习特征和模式,实现高质量的图像重建,在散射介质成像、超分辨成像、低光子成像和无透镜成像等多个场景都取得了令人瞩目的成果[3]。例如,Li等[4]提出了一种可以学习散射介质宏观结构统计特性的CNN模型,实现了稀疏目标散斑图像的散射成像;Dong等[5]提出了端到端的超分辨率卷积神经网络(super-resolution convolutional neural network,SRCNN)可以实时的将低分辨率图像转换为高分辨率图像;Qiao 等[6]提出了基于频域注意力机制的生成对抗网络(deep Fourier channel attention network,DFCAN),实现了低光子条件下的显微成像;Yang等[7]提出了一种能实时校正点扩散函数测量中的系统误差的RNN模型,实现了高质量的无透镜成像。但是,这些方法都仅使用单一模态的数据训练模型,通常泛化能力较差,一旦切换应用场景就必须对模型进行微调,操作较为复杂,并且准确性也会受到限制[8-11]

因此,本文提出一种新型的基于Transformer模块[12]的多模态图像重建模型(multi-modal image reconstruction model based on the Transformer,Trans-MIR)。Trans-MIR使用多个数据源的信息来增强泛化性,扩展应用范围,提高图像重建的质量和效果,并且采用轻量化设计,参数量仅有2 M。在对二维通用人脸散斑图像和STORM采集的三维微管结构超分辨图像进行图像重建的实验中,Trans-MIR表现出高质量图像重建能力和良好的泛化性能。其中,不同景深和散射程度的二维通用人脸散斑图像重建结果的结构相似度(structural similarity,SSIM)指数均高达0.93,三维微管结构超分辨图像重建结果的均方误差(mean squared error,MSE)低至10−4量级。该方法具有进一步扩展至低光子成像、无透镜成像的潜力。

1 图像重建方法 1.1 模型架构

本文提出的Trans-MIR的模型架构如图1所示。该模型采用了编解码器结构,整体呈现U型。模型的前半部分为编码器,可以接受不同模态的数;后半部分为解码器,解码器的最后一层是模型的输出结果。灰色箭头的方向代表数据流动的方向,其中DHWC分别代表数据的维度数、高、宽和通道数。在模型的前向过程中,图像嵌入模块负责将不同维度数的输入图像转换为包含64个通道的特征图,从而将输入数据嵌入到模型之中;Transformer模块能够自适应地调整对特征图中不同信息的关注度,从而高效地提取特征图中所蕴含的全局信息;卷积采样模块使用了残差连接和瓶颈结构,能够实现高性能地特征图采样;隐编码模块对深层特征进行了线性变换,能够强化模型的表征能力;特征压缩模块将64个通道的特征图变换为仅具有前后景2个通道的特征图;最后通过Softmax模块的非线性函数激活得到输出结果。

图 1 Trans-MIR 模型框架示意图 Figure 1 Schematic diagram of Trans-MIR architecture
1.2 Transformer 模块

Trans-MIR中最核心的Transformer模块如图2所示。该模块由批归一化(batch normalization,BN)层、位置编码(positional encoding,PE)层、多头稀疏注意力(multi-head sparse attention,MHSA)机制、线性层和丢弃层等通过级联和残差方式组成。

图 2 Transformer模块示意图 Figure 2 Schematic diagram of Transformer module

BN是一种正则化技术,可以增强模型的非线性建模能力和表示能力。使用BN可以有效地改善梯度消失和爆炸问题,提高模型的训练速度和稳定性,同时降低模型对超参数设置的敏感性。以输入数据中第i个批次的第 k 个特征$ {x}_{ik} $ 为例

$ \begin{array}{c}{B}{N}\left({x}_{ik}\right)={\gamma }_{k}\frac{{x}_{ik}-{\mu }_{k}}{{\sigma }_{k}}+{\beta }_{k} \end{array} $ (1)

式中:$ {\mu }_{k} $$ {\sigma }_{k} $ 分别代表第k个特征在输入数据的全部m个批次中的均值和标准差;$ {\gamma }_{k} $$ {\beta }_{k} $ 是相应的可学习的参数,用于调整输出值的范围和均值,实现重构的操作。

PE是一种为数据增加位置编码的技术。为Transformer模块引入PE可以提供数据中各个位置之间的相对距离信息,以便更好地进行注意力计算和建模。位置编码的数学表达式为

$ \begin{array}{c}{{{\boldsymbol{P}}}{{\boldsymbol{E}}}}_{\left(pos,2i\right)}=\mathrm{sin}\left(\frac{pos}{{10\;000}^{\frac{2i}{{d}_{{\mathrm{m}}}}}}\right)\end{array} $ (2)
$ \begin{array}{c}{{{\boldsymbol{P}}}{{\boldsymbol{E}}}}_{\left(pos,2i+1\right)}=\mathrm{cos}\left(\frac{pos}{{10\;000}^{\frac{2i}{{d}_{{\mathrm{m}}}}}}\right)\end{array} $ (3)

式中:${{{\boldsymbol{P}}}{{\boldsymbol{E}}}}_{\left(pos,2i\right)}$${{{\boldsymbol{P}}}{{\boldsymbol{E}}}}_{\left(pos,2i+1\right)}$ 分别表示输入数据中位置 $ pos $ 的偶数和奇数维度的位置编码向量;$ {d}_{{\mathrm{m}}} $是模型的通道数。

MHSA 是一种基于注意力机制的算法,可以提高模型对于不同特征的抽取能力。针对多模态数据的特点,多头稀疏注意力机制引入了多个头(head),每个头可关注输入数据的不同部分,计算不同的稀疏注意力(sparse attention,SA)分布,并将它们在通道维度组合(concat)起来以获得更加准确和全面的特征表示。计算公式为

$\begin{split} {{MHSA}}\left(X\right)=& {{concat}}({{h}{e}{a}{d}}_{1}\left(X\right),\cdots ,{{h}{e}{a}{d}}_{i}\left(X\right),\cdots,\\ & {{h}{e}{a}{d}}_{N}\left(X\right)) \end{split}$ (4)
$ \begin{split} {{h}{e}{a}{d}}_{i}\left(X\right)=& SA({m}{a}{x}{p}{o}{o}{l}\left(X\right){W}_{i}^{Q}\text{,}{m}{a}{x}{p}{o}{o}{l}\left(X\right){W}_{i}^{K}\text{,}\\ & {m}{a}{x}{p}{o}{o}{l}\left(X\right){W}_{i}^{V}) \end{split} $ (5)
$ \begin{array}{c}{{SA}}\left(Q,K,V\right)={{Softmax}}\left(\dfrac{Q{K}^{{\rm{T}}}}{d}\right)\cdot V\end{array} $ (6)

式中:$ X $ 代表输入的特征数据;${{head}}_{i}$ 代表多头稀疏注意力机制第 i 个头,$ {W}_{i}^{Q}、{W}_{i}^{K}、{W}_{i}^{V} $ 分别是第i个头中的三次线性映射,$ d $ 是一个缩放因子,一般设置为512,Softmax函数将输入映射到 (0,1) 区间内,且所有输出的和为1,对于给定的 n 个输入 $ {x}_{1},{x}_{2},\cdots ,{x}_{n} $ 计算式为

$ \begin{array}{c}{{Softmax}}\left({x}_{i}\right)=\dfrac{{\mathrm{e}}^{{x}_{i}}}{{\displaystyle\sum }_{i=1}^{n}{\mathrm{e}}^{{x}_{i}}},\;\;i=\mathrm{1,2},\cdots ,n\end{array} $ (7)
1.3 数据设置

图3为获取二维通用人脸散斑图像的散射成像系统。激光器(Thorlabs HNL210L,波长为632.8 nm)发出的激光照射到加载了Face-LWF通用人脸图像数据集[13]的相位型空间光调制器(Thorlabs EXULUS-HD2,像素尺寸为8 μm,1 920$ \times $1 200)上进行调制。通过CMOS相机(Thorlabs DCC645C,像素尺寸为3.6 μm,1 280 $ \times $ 1 024)依次采集透过4种不同粒度的毛玻璃散射介质(Thorlabs,DG10-120-MD,125 μm;DG10-220-MD,70 μm;DG10-600-MD,25 μm;DG10-1500-MD,10 μm)的人脸散斑图像。

图 3 人脸散斑数据采集系统 Figure 3 Facial speckle data acquisition system

图3(b)所示,从Face-LWF通用人脸图像数据集中随机选取1 500张不同的人脸图像,在空间光调制器上进行调制,并通过组合4种不同粒度(125 μm,70 μm,25 μm,10 μm)的毛玻璃散射介质与3种不同的CMOS相机位置(离焦面0,20和40 mm)采集了18 000组散斑图像,选取CMOS相机中心区域的800 $ \times $800 像素作为有效数据,并按照6∶2∶2划分训练、验证和测试数据集。

为了获取三维微管结构图像,本文使用SMLM挑战赛2016数据集[14]。将微管宽场图像与STORM超分辨图像组成的三维超分辨数据对进行缩放对齐、随机裁剪、旋转变换,共得到1 400组尺寸为4$ \times $64$ \times $64的有效数据,并按照6∶2∶2划分训练、验证和测试数据集。

1.4 训练细节

为了高效地训练Trans-MIR,并使其收敛到最优,本文使用双线性插值算法将二维通用人脸散斑图像采样为200$ \times $200,除以图像最大灰度值255,将其归一化到0~1,并将三维微管结构图像4个平面的灰度值分别映射到原本灰度值的1%~99%。使用自适应动量估算(adaptive moment estimation,Adam)作为训练优化器,使用交叉熵函数(cross-entropy,CE)作为训练损失函数。模型在NVIDIA GeForce RTX 3090图形计算卡上训练了100个周期,每个周期输入2个批次的图像数据,前50个周期的学习率固定为$ {10}^{-5} $,后50个周期的学习率通过余弦退火算法下降为$ {10}^{-6} $

1.5 评估指标

本文使用SSIM和MSE来评估Trans-MIR的性能。SSIM通过计算2个图像xy的亮度、对比度、结构的差异来表示相似度,计算式

$ \begin{array}{c}{{SSIM}}\left(x,y\right)=l{\left(x,y\right)}^{a}c{\left(x,y\right)}^{b}s{\left(x,y\right)}^{c}\end{array} $ (8)

式中:$ a,b,c $ 为常数;根据实际应用场景确定$ l\left(x,y\right),c(x,y) $$ s(x,y) $ 分别表示亮度相似度、对比度相似度和结构相似度。

MSE通过计算2个图像xy的强度误差来表示相似度,计算式为

$ \begin{array}{c}{{MSE}}\left(x,y\right)=\dfrac{1}{mn}{\displaystyle\sum }_{i}^{m}{\displaystyle\sum }_{j}^{n}{\left[x\left(i,j\right)-y\left(i,j\right)\right]}^{2}\end{array} $ (9)

式中:$ x $$ y $ 分别表示2个图像;$ m $$ n $ 分别表示图像的高和宽;$ x\left(i,j\right) $$ x\left(i,j\right) $ 分别表示两个图像在位置 $ \left(i,j\right) $ 处的像素值。

2 实验结果与分析

为了客观评价Trans-MIR的真实表现,本文将2种模态的测试集数据分别输入到训练好的模型中进行图像重建。图4展示了CMOS相机在离焦面40 mm处采集的二维通用人脸散斑图像的重建结果,图像重建结果中白色代表阳性,黑色代表阴性,绿色代表假阳性,红色代表假阴性。

图 4 不同散射程度二维通用人脸散斑图像的重建结果 Figure 4 Reconstruction results of two-dimensional universal face speckle images with different scattering levels

可以看出,Trans-MIR对不同散射程度的通用人脸图像都能实现高质量的图像重建,重建图像和前景、背景都相当干净,仅在一些复杂纹理的边缘细节上出现微小的误差。从整体上看,重建图像与真实图像的眼睛、鼻子、嘴巴、耳朵和头发等图像细节和纹理几乎一致。这充分说明了Trans-MIR具有不错的图像重建能力和泛化能力。

表1给出了人脸图像在COMS相机距离焦平面不同位置的重建结果。分析表1中的数据可知,对于CMOS相机在离焦面0 mm处采集的二维通用人脸图像,Trans-MIR重建结果的SSIM达到0.942,远好于Guo等[15]在相似任务上约0.8的重建结果。并且可以发现,该模型对于CMOS相机在不同位置下采集的二维通用人脸图像都能高质量的重建,SSIM均高于0.93,进一步说明了Trans-MIR具有很好的图像重建能力和泛化能力。

表 1 二维通用人脸图像的重建结果 Table 1 Reconstruction results of two-dimensional universal face images

Trans-MIR对三维微管结构超分辨图像的重建结果准确度同样很高,在测试数据集的280组图像上的平均MSE为5.8×10−4 ,达到了10−4 量级,在相似的任务上DFGAN的表现仅为 10−3 量级。图5为一组三维微管结构超分辨图像的重建结果,其中D1,D2,D3,D4 分别代表三维数据在不同平面的二维切片。

图 5 三维微管结构图像的重建结果 Figure 5 Reconstruction results of the three-dimensional microtubule structure image

通过观察可以发现,Trans-MIR对三维微管结构图像的各个平面的重建结果都很好,相较于原始模糊的输入图像,重建结果的清晰度得到大幅提高,已经可以准确地分辨微管结构,并且得益于多模态模型,重建结果抑制了部分在真实图像中微管结构末尾处的噪声信号。实验结果充分表明,Trans-MIR同样可以实现三维微管结构图像的高质量重建。

3 结 论

在计算成像领域,当前大多数的图像重建方法都集中在针对单一模态的数据上,而对于多模态图像重建方向的研究则相对较少。本文提出一种基于Transformer模块的多模态的图像重建模型Trans-MIR。Trans-MIR使用二维通用人脸散斑图像和三维微管结构超分辨图像2种模态的数据进行训练,训练好的Trans-MIR实现了对多模态的图像数据的高质量图像重建。Trans-MIR在模型架构上做出了一定的改进,使用轻量化设计的Transformer模块匹配多模态图像重建任务,对于研究多模态图像重建有一定的启发作用。Trans-MIR可以进一步扩展至低光子成像、无透镜成像场景。

参考文献
[1] 左超, 陈钱. 计算光学成像: 何来, 何处, 何去, 何从?[J]. 红外与激光工程, 2022, 51(2): 20220110.
[2] 邵晓鹏, 刘飞, 李伟, 等. 计算成像技术及应用最新进展[J]. 激光与光电子学进展, 2020, 57(2): 020001.
[3] BARBASTATHIS G, OZCAN A, SITU G. On the use of deep learning for computational imaging[J]. Optica, 2019, 6(8): 921–943.
[4] LI S, DENG M, LEE J, et al. Imaging through glass diffusers using densely connected convolutional networks[J]. Optica, 2018, 5(7): 803–813.
[5] DONG C, LOY C C, TANG X O. Accelerating the super-resolution convolutional neural network[C]//Proceedings of the 14th European Conference on Computer Vision. Amsterdam, The Netherlands: Springer, 2016: 391 – 407.
[6] QIAO C, LI D, GUO Y T, et al. Evaluation and development of deep neural networks for image super-resolution in optical microscopy[J]. Nature Methods, 2021, 18(2): 194–202.
[7] YANG J Y, YIN X J, ZHANG M X, et al. Learning image formation and regularization in unrolling AMP for lensless image reconstruction[J]. IEEE Transactions on Computational Imaging, 2022, 8: 479–489.
[8] KNOLL F, HOLLER M, KOESTERS T, et al. Joint MR-PET reconstruction using a multi-channel image regularizer[J]. IEEE Transactions on Medical Imaging, 2017, 36(1): 1–16.
[9] BOUSSE A, BERTOLLI O, ATKINSON D, et al. Maximum-likelihood joint image reconstruction/motion estimation in attenuation-corrected respiratory gated PET/CT using a single attenuation map[J]. IEEE Transactions on Medical Imaging, 2016, 35(1): 217–228.
[10] ZANFIR A, ZANFIR M, GORBAN A, et al. HUM3DIL: semi-supervised multi-modal 3D HumanPose estimation for autonomous driving[C]// Proceedings of the 6th Conference on Robot Learning. Auckland: PMLR, 2023: 1114 – 1124.
[11] AN S Z, LI Y, OGRAS U. mRI: multi-modal 3D human pose estimation dataset using mmWave, RGB-D, and inertial sensors [EB/OL]. [2023–03–01]. https://openreview.net/pdf?id=Oa2-cdfBxun.
[12] VASWANI A, SHAZEER N, PARMAR N, et al. Attention is all you need[C]//Proceedings of the 31st International Conference on Neural Information Processing Systems. Long Beach: Curran Associates Inc. , 2017: 6000 – 6010.
[13] HUANG G B, MATTAR M, BERG T, et al. Labeled faces in the wild: a database forstudying face recognition in unconstrained environments[C]//Proceedings of the Workshop on Faces in 'Real-Life' Images: Detection, Alignment, and Recognition. Marseille: Erik Learned-Miller and Andras Ferencz and Frédéric, 2008.
[14] SPEISER A, MÜLLER L R, HOESS P, et al. Deep learning enables fast and dense single-molecule localization with high accuracy[J]. Nature Methods, 2021, 18(9): 1082–1090.
[15] GUO E L, ZHU S, SUN Y, et al. Learning-based method to reconstruct complex targets through scattering medium beyond the memory effect[J]. Optics Express, 2020, 28(2): 2433–2446.