日前,美图影像研究院(MT Lab)联合北京交通大学提出了超高分辨率抠图方法MEMatte (Memory Efficient),一个显存友好的自然图像抠图框架用于在显存受限的场景下(如商用显卡与边缘设备)进行高清图像的精细抠图。文章已被人工智能顶会 AAAI 2025 接收。
随着摄影设备的性能升级,我们在日常采集到的图像分辨率已经能达到4K乃至8K,这种高分辨率图像有着更好的视觉质感与更丰富的细节内容。而自然图像抠图方法相比于一般的分割方法更善于区分细节充实的场景,因此适用于更复杂的下游任务诸如图像与视频的后期编辑。
当前Transformer (ViT) 为自然图像抠图方法带来了显著的性能提升,这种提升来自于全局注意力模块对图像内容的长程建模能力。然而,基于ViT的抠图方法在计算资源受限的场景下无法处理高分辨率图像,因为ViT内部的全局注意力会产生次方倍的计算开销。此外,抠图任务的特性要求输入能同时保留清晰图像中的细节信息与完整图像中的语义信息,因此无法采用降采样或图像切片等手段来降低模型输入的图像尺寸。这些限制条件要求我们在维持图像输入大小的前提下降低全局注意力对计算资源的消耗。
过往的解决方案包括令牌剪枝与令牌融合,这些方法都需要在模型的不同层以固定比例逐渐丢弃冗余令牌。但是这类方法存在两种缺陷:
(1)丢弃令牌会导致抠图效果显著变差,因为被丢弃的令牌中存在着需要被保留的细节信息;
(2)预设固定的丢弃比例无法适应不同的输入,对于复杂样例丢弃过多信息会导致精度下降,而对于简单样例则依然存在计算冗余。
为了解决上述问题,研究团队提出的MEMatte采用了双分支的令牌路由设计。该设计的核心思想是将包含语义信息的令牌送入全局注意力模块进行处理,并且将其余令牌送入精心设计的轻量化令牌提炼模块 (Light-weight Token Refinement Module, LTRM) 进行处理。令牌的分流是通过动态路由机制 (Batch-constrained Adaptive Token Routing, BATR)进行筛选,在推理时无需预设固定的比例,因此能够输入自适应地进行调整。这一设计显著降低了模型的计算开销,并在Nvidia GeForce 2080Ti 商用 GPU上实现了4K分辨率图像抠图。此外,研究团队还开源了超高分辨率自然图像抠图数据集UHR-395 (Ultra High Resolution dataset),用于高分辨率模型的训练与评估。
图1:MEMatte方法与已有方法在不同分辨率的显存消耗对比对比
图2:UHR-395数据集与其它开源抠图数据集的输入下UHR-395有着最高的平均分辨率。
方法
在推理过程中,MEMatte的ViT编码器内部的每一个全局注意力模块前都存在一个路由器 (Router)。这个路由器通过局部-全局策略为全局注意力进行令牌重要性评估。该策略对每层的图像输入进行重新映射以及通道拆分后的全局池化操作来兼顾局部与全局信息。此后通过LogSoftmax (LS)操作形成令牌的分流结果。
Router 将分流到全局注意力之外的令牌送入 LTRM 进行处理。LTRM 由数个轻量化组建构成,其中包括两组映射层。此外,深度卷积 (Depth-Wise Convolution, DWC) 层处理局部空间信息,高效通道注意力层 (Efficient Channel Attention, ECA) 处理全局特征信息。
在训练MEMatte时，研究团队采用 BATR 的训练机制，搭配蒸馏学习策略与常规抠图训练损失，实现快速且平稳的模型训练，在训练过程中，假定从数据分布中均匀地采样，则每次的压缩比例应当为一个定值。基于这项假设， BATR 为模型预设了分流比例
来约束模型预测的批次平均分流比例，定义批次平均分流比例 
为如下形式：
通过约束模型的预测值与预设值的差异,训练出足够有效的 Router 以自适应的处理各种图像。
图3:MEMatte的训练与推理流程
试验结果
实验结果表明,MEMatte在现有的基准以及UHR-395的高分辨率测试集上都领先于过去的方法,并在域外测试集上展现了良好的泛化性能。此外,在抠图测试集Composition-1K上,MEMatte相比于其baseline节省了约88%的显存开销,并降低了约50%的推理时间。
图4:MEMatte与过往方法在通用抠图基准上的比较
图5:MEMatte与过往方法在UHR-395高分辨率测试集上的比较
图6:在高分辨率图像上的抠图效果。
研究图案对 Router 分流结果的可视化结果展示了MEMatte在不同的全局注意力层对令牌的选择倾向。MeMatte使中间的自注意力层主要关注物体边缘与细区域的复杂区域,同时避免了令牌通过ViT浅层的自注意力操作,这大幅提高了输入自注意力层的令牌稀疏性。
图7:对Router分流结果的可视化。其中灰色区域是分流向 LTRM 的令牌。
超高分辨率自然图像抠图数据集
目前公开的抠图数据集难以作为高分辨率抠图的基准,因为它们无法同时兼顾类别多样性和高分辨率。为了解决这一问题,研究团队提出了一个超高分辨率自然抠图数据集UHR-395。该数据集包含11个类别,共395个前景物体,划分为训练集355张以及测试集40张,平均分辨率高达4872×6017。
图8:前景类别分布。
图9:图像分辨率分布。
为确保标注数据的高质量,每个前景物体的标注结果都通过多轮严格的审核机制,并由抠图领域专家进行最后确认,不合格标注结果均被返回进行重新修正,最终获得高质量的自然图像抠图标注数据。
图10:标注结果展示。
结语
超高分辨率的自然图像抠图方法对实现复杂图像编辑与视频后期处理等下游任务而言至关重要,在本研究中,来自美图影像研究院(MT Lab)和北京交通大学的研究人员共同提出了 MEMatte 和超高分辨率数据集 UHR-395。与现有的抠图方法相比,采用MEMatte 能够显著降低推理过程中的内存使用量和延迟,在高分辨率和真实世界数据集上实现了最先进的性能,突出了令牌路由机制 BATR和 轻量化令牌细化模块 LTRM 的稳健性,为美图的图像与视频处理能力提供高性能且灵活的技术支撑,为用户带来高效优质使用体验。此外, MEMatte 也有望促进高分辨率抠图的进一步研究。
