51c视觉~合集8

#Mixture-of-Experts, MoE

ViT的轻量化还能做到什么程度？MoE方法立大功！使用专家混合（Mixture-of-Experts, MoE）来简化而非增强视觉 Transformer 。

在本文中，作者探讨了一种策略，该策略使用专家混合（Mixture-of-Experts, MoE）来简化而非增强视觉 Transformer 。MoE层中的每个专家都是一个SwiGLU前馈网络，不采用复杂的注意力或卷积机制。

逐深度缩放被应用于逐步减少隐藏层的大小，并且分阶段增加专家的数量。使用了分组 Query 注意力。作者研究了在小数据集上进行预训练和不进行预训练的所提方法，并探讨了在这种规模下迁移学习是否有效。作者发现，即使参数量仅有0.67M，该架构也具有竞争力。

1 Introduction

在计算机视觉的实际应用中，例如边缘智能，小型且性能高效的模型仍然被优先考虑以克服计算挑战[1]。视觉 Transformer （ViTs）[2]取得了显著成果，但当模型大小和数据集较小时，其性能会大幅下降[3]。因此，有研究开始探究在中等规模数据集上表现良好的轻量级视觉 Transformer 。几乎所有这些研究都使用新型注意力块[4, 5]或将卷积机制[6]整合到其架构中。

另一方面，Tan[7]通过使用 Mask 自动编码器（MAE）[8]作为预训练策略，证明了可以使ViT有效地从小数据集中学习。在该研究中，ViT由12个 Transformer 编码层组成，每个编码层包含一个多头注意力组件和一个前馈网络。前馈网络由两个线性层组成：第一个将输出扩展到嵌入大小的两倍，而不是四倍，第二个将输出减小回嵌入大小。为了进一步减轻模型负担，减少前馈网络中间的扩展输出大小可能会有帮助，但过度减少可能会负面影响模型性能。

考虑到这些因素，作者设计了一种使用专家混合（MoE）[9]来简化视觉 Transformer 的架构。在作者的架构中，每个MoE层中的每个专家由一个SwiGLU[10]前馈网络形成。从设计上讲，SwiGLU在参数方面比典型的多层感知器要重。然而，通过在MoE层中设置多个专家，作者能够使SwiGLU中的隐藏大小小于嵌入大小，而不会对模型性能产生负面影响。此外，作者在每个SwiGLU中将3个线性变换中的2个在层间共享。这有助于显著降低参数数量，同时保持MoE的优势。除此之外，为了进一步减少参数数量，作者按照深度方向逐步增加MoE中的专家数量，同时线性地减小隐藏大小，借鉴了深度方向缩放[11]的思想。最后，作者使用分组 Query 注意力[12]以保持参数数量低。源代码将在不久的将来提供。

2 Method

作者提出的方法包括两个部分：mLIT 和 mmLIT。mLIT 是一种基于专家混合的轻量级视觉 Transformer ，而 mmLIT 是通过 Mask 自动编码器预训练策略进行预训练和微调的 mLIT。

mLIT

与原始的ViT[2]类似，作者首先将每个图像划分为不重叠的块。这些块通过线性变换转化为一系列嵌入向量，并通过可学习的位置嵌入进行增强。处理后的嵌入向量经过一系列基于MoE的 Transformer 编码器层。图1展示了作者 Transformer 编码器层的整体结构。

2.1.1 Grouped Query Attention (GQA)

分组 Query 注意力将 Query 头分为个组，每个组共享一个键头和一个值头。例如，GQA-1，只有一个组，等同于多 Query 注意力。

2.1.2 专家混合（MoE）

MoE的损失函数。为了鼓励所有专家具有同等的重要性，引入了两种损失函数：负载均衡损失和重要性损失。为了计算负载均衡损失，作者首先通过以下概率确定

在mLIT中，每个专家中的和及其对应偏置在MoE层中是共享的。

2.1.3 Depth-wise Scaling

作者的视觉 Transformer 在 Transformer 编码器层之间并没有固定的隐藏尺寸。相反，隐藏尺寸是从第一层到最后一层线性减少的。

此外，mLIT中的专家数量在不同阶段会增加，具体来说，在每3、4或5层之后，从3增加到5，这取决于模型的大小。更多详细信息请参见图4。

mmLIT

在mmLIT中，作者将在mLIT上应用 Mask 自动编码器。作者严格遵循文献[7]中概述的流程，该流程基于原始MAE论文[8]。Mask 比例设定为0.75。然而，与文献[7]不同的是，作者没有在解码器的输入中合并可分离的学习位置嵌入（见图5）。与文献[7]类似，作者计算自动编码器的损失为 Mask 块上的损失与未 Mask 块上的额外折减损失的之和。总损失由以下公式给出：

3 Experimental Setup

作者研究了mLIT和mmLIT在三种不同规模：S、XS和XXS上的性能表现。表1展示了每个模型的编码器细节，表2展示了解码器细节。对于mmLIT，作者在Cifar100 [14] 训练数据集上进行了自监督预训练，并在Cifar100、Cifar10[14]、Flowers102 [15]和Svhn [16]上进行微调。此外，作者还对上述四个数据集进行了有监督学习（不进行任何预训练）。

图3：MoE的实际实现。与输入仅是一个向量时的情况不同，当个嵌入进入层时，几乎所有的专家网络都会接收到一些由分配器发送的嵌入。

输入模型的图像大小为36 x 36，略大于这些数据集原始的32 x 32尺寸（除了Flowers102，其尺寸更大且多样）。这导致每个图像被划分为144个 Patch ，给定3 x 3的 Patch 大小。与[7]类似，在预训练和微调阶段，每张图像都添加了一个辅助的虚拟 Patch 。这个虚拟 Patch 对所有元素的值均为零，作为第145个 Patch 用于分类目的。

基于MoE的 Transformer 编码层中的所有线性层均包含偏置。然而，作者从编码器和解码器中的其他线性投影层中排除了偏置。对于所有层类型的权重和偏置初始化，作者依赖于Pytorch提供的默认方法。对于层归一化[17]的处理方式也同样如此，作者使用Pytorch的默认设置。

Pre-training

mmLIT-S、mmLIT-XS 和 mmLIT-XS 分别在 Cifar100 数据集上预训练了4000、6000和8000个周期。作者使用了 AdamW 优化器 [18]，并将权重衰减设置为 0.05。总周期的前5%用于 Warm up  [19]。随后，作者采用了余弦衰减计划 [20] 来调整学习率，并遵循线性学习率缩放规则，基础学习率为：

微调

对于每个预训练模型，作者进行了两组微调实验。首先，作者对在 Cifar100 上预训练的模型进行了针对 Cifar100 分类任务的微调，共300个周期。其次，作者通过在 Cifar100、Cifar10、Flowers102 和 SVHN 上对每个模型进行100个周期的微调，来评估模型的迁移学习能力。表4展示了第一组微调的配置，而表5则展示了第二组微调相对于表4的配置差异。

此外，作者还评估了不同尺寸的 mLIT 模型（仅使用监督学习从头开始训练）的性能。表6展示了在四个数据集上对 mLIT 进行监督学习时的配置差异。作者遵循线性学习率缩放规则（方程16）进行所有微调和监督学习。所有预训练和微调都使用混合精度（torch.float16）进行。

4 Result

本研究的结果部分旨在详尽地展示实验数据和模拟结果的详细信息。以下各节将分别阐述所采用的方法在不同测试集上的表现，并对所得结果进行分析讨论。

On Pre-Training

图6展示了不同规模的mmLIT预训练损失。当训练损失在x轴（时期）和y轴（损失值）上以对数尺度绘制时，损失下降表现出近似对数线性关系。这表明损失在每个时期呈指数级减少。

mnLIT-S显示出最低的损失，在第400个时期达到0.020404。mmLIT-XS和mmLIT-XXS在第6000和第8000个时期分别实现了0.021051和0.021514的损失。mmLIT-S的预训练稳定性显著，其训练损失的变化最小。相比之下，mmLIT-XS和mmLIT-XXS的训练损失显示出更大的波动性和更频繁、更明显的高峰，特别是在第1000个时期之后尤为明显。

On Fine-Tuning

表7比较了在300个训练或微调周期后，具有相似大小的模型的性能。mmLIT-S的准确率比Mac-VIT-C100高近1%，尽管其参数仅有Mac-VIT-C100的三分之二。mmLIT-XS的大小是Mac-VIT-C100的三分之一，其性能仅低1.5%。更有趣的是，仅有Mac-VIT-C10018%参数的mmLIT-XXS依然具有竞争力，准确率仅低3.3%。此外，mmLIT-XXS甚至略微超过了ResNet56，尽管ResNet56的参数多出18%。

表8展示了仅用Cifar100数据集进行预训练的模型的迁移学习能力，与文献[21]报告的结果进行对比。可以看出，即使参数规模为0.67M的mmLIT也具有竞争力，此时VIT-t+SSAT的规模几乎是mmLIT-XXS的9倍，而其余模型至少是mmLIT-XXS的30倍以上。此外，除了在架构中包含卷积的CVT-13之外，mLIT在Cifar100、Cifar10和Svhn数据集上，即使在小规模下也始终优于普通的VIT和Swin-T。Flowers102是一个细粒度分类数据集，因此mLIT在36 x 36的图像大小上不具备竞争力并不意外。

5 讨论

表8的结果表明，在更小的规模下，ViT可以通过简化的MoE进行更好的学习。在作者当前的 Proposal 中，在MoE层中共享。这种安排背后的主要考虑是上应用的非线性sliu。尽管如此，作者还探索了mmLIT-S上的另外两种可能的共享安排（见表9），结果表明保留会带来轻微的优势。

在作者的设计中，作者逐步减少了隐藏层的尺寸并在几个阶段增加了专家的数量。这种安排受到卷积神经网络[27]的启发，在卷积神经网络中，特征图的大小会逐层减少，而随着网络的加深，特征图的数量会增加。作者发现，如果早期层的减少不显著，且在最后几层不超过75%，MoE的这种安排下隐藏层尺寸的减少对性能的影响很小。未来的实验需要进一步确认这一点。

当作者在四个数据集上进行微调时，作者首先将用于Cifar100的微调设置应用于其他三个数据集。但是，早期阶段的过拟合成为一个一致性问题，因此作者必须修改设置以适应每个数据集。当作者尝试从头开始使用mLIT进行有监督学习时，也遇到了类似的问题。

因此，为了清晰起见，作者在表4和表5中报告了设置。值得注意的是，作者没有对配置进行彻底搜索；

作者只是调整了超参数以避免训练过程中的过拟合。此外，与[21]不同，作者在所有微调中并未使用先进的增强技术。

在mmLIT的预训练期间，作者针对不同尺寸的模型使用了不同数量的训练轮次。从图6可以看出，从损失的角度来看，应该为mmLIT-XS和mmLIT-XXS分配更多的轮次。然而， Mask 自编码器的一个缺点是模型的损失与其在后续下游任务中的性能之间的关系不明确。根据作者的经验，在后期阶段再进行几百轮的预训练，有时会导致分类性能相似或更差。这可能是在数据集较小时的特定问题。尽管如此，作者相信对于这类数据集，进行数千轮的进一步预训练应该能带来更好的模型。

在本研究中，作者发现仅用50,000张图像进行预训练的mLIT即使在最小尺寸下也能作为某种基础模型[28]。在Cifar10和Flowers102上，通过Cifar100进行预训练的模型至少提高了10%。在Cifar10上，mmLIT-XXS在不到40轮次的情况下可以达到90%的准确率。在Svhn上，改进是适度的。这部分是因为mLIT即使没有预训练也表现良好。也可能是由于Cifar100中缺少类似的图像。作者相信，在稍大一些且更多样化的数据集上进行预训练的模型能够在这些微小尺寸上的各种简单任务上具有竞争力的表现。

在文献中，通常会将卷积层[29, 23]集成进来，或者采用卷积“教师”模型[24]来赋予归纳偏置。似乎使用简化的MoE（多专家系统）可以帮助缓解归纳偏置的不足。然而，正如[7, 21]所示，通过 Mask 自动编码器的设置，可以克服归纳偏置的问题。

6 Conclusion

在本文中，作者展示了精简的MoE架构在创建极为轻量级的视觉 Transformer 方面的潜力。

通过在MoE层间共享参数并采用深度可分离缩放，作者实现了模型复杂度与性能之间的平衡。

此外，作者的研究结果表明，在稍大且更多样化的数据集上进行预训练可以增强模型在各种任务中的通用性和有效性。

精简的MoE方法在与 Mask 自动编码器设置结合使用时，似乎在缓解归纳偏置不足的问题上具有很大前景。

#TAPTR

让跟踪最细粒度的视觉元素像检测一样简单

本文介绍了两篇关于视觉元素跟踪的系列工作，均收录于ECCV 2024。本文提出了一种新的网络架构TAPTR，将跟踪任务视为检测任务，通过将每个跟踪点视为一个visual prompt，并使用point-query和cost-volume等技术进行跟踪。

本文介绍一下我们最近的两篇系列工作:

TAPTR: Tracking Any Point with Transformers as Detection（收录ECCV 2024）

TAPTRv2: Attention-based Position Update Improves Tracking Any Point

主页：https://taptr.github.io

代码：github.com/IDEA-Research/TAPTR

TAPTR：https://arxiv.org/abs/2403.13042

TAPTRv2：https://arxiv.org/abs/2407.16291

TAPTR（收录ECCV 2024）

图1. TAPTR的网络架构图。

如图1所示，我们首先提出了TAPTR这个简单高效的网络架构。与前人的工作不同，我们不再将每个被追踪点的各种特征 (包括语义特征、光流、光流的映射、可见性、cost-volume等) 混为一谈，然后一股脑地送入到网络中进行处理，使得对点追踪的建模黑盒化。在TAPTR中，我们首次提出将每个要被追踪的点都清晰地看作是point-level的visual prompt，而TAP任务则可以被看作是从视频的每一帧中将这个visual-prompt所指代的点给检测出来。这使得我们对追踪点的建模十分清晰，同时基于前人对检测任务的详细研究也使得整个架构简单且高效。TAPTR可以被大致地分为三个阶段，接下来我会按次序进行讲解。

第一个阶段为准备阶段。首先是对video中每一个frame的特征图进行提取。这里我们使用的是一个ResNet以及一个Transformer Encoder。在这之后，我们需要为每一个被追踪点获取一个能够准确描述该点的特征。为了方便描述，不失一般性的，我们在这里假设我们只需要追踪一个点。在TAPTR中，我们会在被追踪点开始被追踪的位置l_e进行双线性插值从而采样出描述该点的特征f_e。为了将被追踪点从每一帧中检测出来，我们需要在每一帧 都构建属于该被追踪点且负责将该点检测出来的point-query。我们将这些point-query初始化为 (l_e, f_e).

由于tracking-any-point任务与optical flow estimation任务的相似性，我们将在optical flow以及前人工作中被验证十分有效的cost-volume (有些工作也称其为correlation map) 也放入到我们的framework中。为了构建每一个点的cost-volume，我们将每一个被追踪点的初始特征f_e与每一帧的特征图进行点乘。（由此可见，随着被追踪点的数量以及输入图像的增大，这里的cost将会急速增高，这也是我们v2版本致力于解决的问题之一）

第二个阶段是通过我们的point-decoder将被追踪点在各个帧中检测出来。在point-decoder中，主要包含四个部分。由于cost-volume提供了最基础的被追踪点与图像的相似度关系，为被追踪点的检测提供了最基本的感知，我们首先会通过cost-volume aggregation模块将point query和cost-volume进行交互。接着，我们也会将point query与视频图像的深度特征图之间通过cross-attention进行交互从而补充更加详细的几何结构信息。接着，在同一帧内的所有point query之间会互相交互，从而扩大各个point query的感受野，提供更加丰富的上下文信息。此外，我们会将属于同一个被追踪点的在各个frame下的point query在temporal-attention中互相进行交互，从而补充时序信息。每一个point query最终都会经过content updater以及position updater来更新它的语义特征、位置以及可见性。

第三阶段：由于一个视频通常都很长，为了节省计算资源，同时也是为了适配不同长度的视频，在TAPTR中我们使用滑动窗口的策略。在每一次forward的时候，我们仅仅需要将所有点在某一个window内的视频帧中检测出来即可。检测结果将会通过window post-processing模块更新到完整的trajectory中。当滑动窗口滑动过整个视频之后，我们的追踪任务也就完成了。

表1. TAPTR与其他方法的性能对比。

我们在TAP-Vid benchmark上对TAPTR的性能进行了测试，结果如表1所示。可以看到，TAPTR在各个数据集中都达到了最优的性能，以及最快的速度。

表2. TAPTR的消融实验。

由于TAPTR是TAP任务的baseline工作，如表2所示，我们提供了丰富的ablation实验，验证了TAPTR中各个模块的有效性来作为后续工作的参考，从而方便后续的研究工作。

TAPTRv2

正如我们上文对TAPTR的讲述中所提到的那样。受optical flow的影响，现有的方法，包括TAPTR，都十分依赖cost-volume。这样的依赖不仅仅会导致额外的计算量（尤其是在被追踪点数以及视频的分辨率较大时），还会污染我们的point-query的语义特征。

但是，cost-volume却对TAPTR的性能有着很大的影响。因此，在TAPTRv2中，我们主要探究了：1. cost-volume为什么work？2. 能否有更加轻量化地方式来得到cost-volume或者它的等价替代者？3. cost-volume或者它的代替者要怎么被使用才能不污染我们的point-query？

表3. cost-volume对domain gap的缓解作用的实验证明。

如表3所示，我们发现cost-volume真正起作用的原因可能是在于它对domain gap的缓解作用。

图2. 在APU帮助下的TAPTRv2的decoder层。

我们认为，cost-volume与注意力机制中的注意力有着本质的相似性。因此我们可以用attention value来作为cost-volume的替代者。同时，为了防止其污染point-query，我们不再简单地将其融入到point-query的语义特征中，而是将其直接转化成一个对point-query的位置的一次更新。为此，如图2所示，我们提出了Attention-based Position Update (APU) 。

图3. TAPTRv2的架构图。

图4. TAPTRv2、TAPTR以及前人的工作的方法的对比。

如图3以及图4所示。在APU的帮助下，我们不仅仅能够在不污染point-query的情况下保留原本的domain gap的缓解能力，同时，我们也无需再额外计算cost-volume，也无需在decoder中加入额外的cost-volume aggregation操作了。整个模型十分的简洁高效。

表4. TAPTRv2与其他方法在性能上的对比。

此外，如表4所示，在保持了对domain gap的缓解能力并解决了point-query被污染的问题之后，TAPTRv2相比于TAPTR在各项指标上有了进一步的提升。

#DIVA

Diffusion反馈强势助力CLIP秒变火眼金睛：北京智源研究院、中科院自动化所联合推出DIVA

本文分享论文Diffusion Feedback Helps CLIP See Better，专注于通过自监督学习范式解决CLIP无法区分细粒度视觉细节的问题。

• 作者：王文轩（中科院自动化所-智源研究院联培博一研究生），孙泉（智源研究院视觉模型研究中心算法研究员），张帆（智源研究院视觉模型研究中心算法研究员），唐业鹏（北交博一研究生），刘静（中科院自动化所研究员），王鑫龙（智源研究院视觉模型研究中心负责人）
• 单位：中科院自动化所，中国科学院大学，北京智源人工智能研究院，北京交通大学
• 论文链接：https://arxiv.org/abs/2407.20171
• 项目主页：https://rubics-xuan.github.io/DIVA/
• 相关代码链接：https://github.com/baaivision/DIVA

动机何在？——CLIP视觉缺陷

对比语言-图像预训练（CLIP）在跨领域和跨模态的开放世界表示方面表现出色，已成为各种视觉和多模态任务的基础。

自从CLIP被提出以来，近年来涌现了许多关于CLIP模型的后续研究。这些研究通过预训练和微调CLIP模型，取得了性能提升并开发了新的能力。然而，这些方法仍然存在不可避免的局限性，因为它们高度依赖于图像-文本数据对，无法仅在图像数据上实现预期效果。

此外，最近的不少研究指出，尽管CLIP在零样本任务中表现出色，但由于对比学习范式和训练中使用的噪声图像-文本对，其在感知理解方面存在一些局限性。这些局限性包括难以准确理解长文本和难以辨别相似图像中的细微差异。虽然一些研究试图解决长文本理解问题，但改善CLIP的细粒度视觉感知能力的研究仍然不足。感知视觉细节的能力对于基础模型至关重要，而CLIP在这方面的不足直接影响了以CLIP作为视觉编码器的视觉和多模态模型的表现。

因此，在这项工作中，我们专注于通过自监督学习范式解决CLIP无法区分细粒度视觉细节的问题。基于文本到图像的扩散模型能够生成具有丰富细节逼真图像的先验，我们探索了利用扩散模型的生成反馈来优化CLIP表征的潜力。

如何解决？——Diffusion Feedback来优化CLIP视觉细节表征

我们提出了一种简单的CLIP模型后训练方法，通过自监督扩散过程在很大程度上克服了其视觉缺陷。通过使用CLIP的密集视觉特征对扩散模型进行条件化，并将重建损失应用于CLIP优化，我们将扩散模型作为CLIP的视觉助手，因此我们将该框架命名为DIVA。

具体而言，如图2所示，DIVA主要由两个部分组成：一是需要增强视觉感知能力的CLIP模型，二是提供生成反馈的预训练扩散模型。输入原始图像和空文本（图2中标记为"Null"）后，CLIP模型会编码相应的视觉特征，这些特征将与来自扩散模型文本编码器的空文本嵌入结合，为扩散过程提供条件。对于添加了噪声的图像，扩散模型尝试在上述条件下预测从前一步到当前步骤中添加的噪声。在训练过程中，除了CLIP模型外，所有部分的权重都保持不变，训练目标只是最小化重建损失（即扩散反馈指导）。通过这种方式，通过约束扩散模型更准确地预测添加的噪声，CLIP的原始语义丰富的判别表示将通过扩散反馈逐渐优化为包含更多视觉细节的表示。

此外更有意思的是，DIVA不需要额外的文本标注数据，只需可轻易获取的纯图片数据就能大幅使得CLIP弥补其视觉感知短板，这一点相比之前方法收集大量图文数据对的高昂成本是非常难得的！

效果如何？——立竿见影！

为了评估DIVA的有效性并展示其增强CLIP表示的潜力，我们在多模态理解和视觉感知任务上进行了全面的实验。

视觉细粒度感知方面

为了验证DIVA能够有效缓解CLIP模型固有的视觉能力不足，我们首先在各种现有的CLIP模型上进行了实验。DIVA在评估视觉-语言模型视觉能力的MMVP-VLM基准测试中使得现有的多个CLIP模型的性能取得了显著的提升（提高了3-7%）。

作为更强的视觉骨干网络为多模态大模型和视觉模型带来的收益评估

接下来，在DIVA的帮助下，我们进一步评估了增强后的CLIP骨干网络在多模态理解和视觉感知任务中带来的性能提升。DIVA的优势在于它不仅仅能让CLIP变聪明，还能让那些基于CLIP的大型多模态语言模型以及视觉模型变得更加厉害。在这些多模态和纯视觉的基准测试上准确率的显著提升，得益于我们DIVA范式通过生成反馈大幅增强了CLIP的视觉感知能力。

CLIP泛化能力评估

在全面验证了我们的方法提升CLIP模型细粒度视觉感知能力的效果后，我们进行了CLIP模型原始泛化能力的全面评估。在只由纯图片数据驱动整个框架的前提上，DIVA能够保持CLIP原本优秀的泛化性能。29个图片分类和图文检索的基准测试上无论是看图识物还是找图配字的实验结果都能证明，经过DIVA优化视觉表征之后的CLIP模型能够保留CLIP原本优秀的泛化能力。

未来展望？——大有可为！

当前局限

1. 数据和模型规模可进一步扩展。
2. 由于这篇工作只是该方向的一个开始，目前仅展示了生成扩散模型用于提升CLIP模型表示的潜力，当前主要关注设计一个简单但有效的框架。

未来可探讨的方向

1. 可以结合更细粒度的监督方案进一步提升CLIP模型的能力。
2. 扩展超越图像-文本数据的其他模态，如视频和音频。
3. 发展基于扩散模型的更通用、更强大的框架，以增强视觉-语言基础模型。

#Segformer~PCB缺陷检测

 PCB缺陷检测是电子制造的一个重要方面。利用Segformer等先进模型不仅可以提高准确性，还可以大大减少检测时间。传统方法涉及手动检查，无法扩展且容易出错。利用机器学习，特别是 Segformer模型，提供自动化且精确的解决方案。

实现步骤

    下面是具体步骤： 

  【1】安装所需环境。首先，我们安装所需的库。其中，albumentations用于数据增强，transformers允许访问 Segformer等预训练模型，并xmltodict帮助解析数据集的XML注释。

pip install evaluate albumentations transformers accelerate xmltodict

  【2】数据集。这个项目中使用的数据集由Roboflow提供。可以从下面链接获取：

https://universe.roboflow.com/diplom-qz7q6/defects-2q87r/dataset/16

该数据集分为测试文件夹和训练文件夹，由XML格式的图像及其相应注释组成。

# Create train and test sets
train_folder = "drive/..../train/images/"
test_folder = "drive/.../validation/images/"


train_img_paths = sorted([train_folder + f for f in os.listdir(train_folder) if f.endswith("jpg")])
test_img_paths = sorted([test_folder + f for f in os.listdir(test_folder) if f.endswith("jpg")])
train_xml_paths = [f[:-3] + "xml" for f in train_img_paths]
test_xml_paths = [f[:-3] + "xml" for f in test_img_paths]


train_ds = {"image_paths": train_img_paths, "xml_paths": train_xml_paths}
test_ds = {"image_paths": test_img_paths, "xml_paths": test_xml_paths}

    XML文件包含多边形标注信息，指示PCB图像上缺陷的位置。该函数process_mask读取XML标注信息并将其转换为掩码（类似图像的数组）。该掩模对应于PCB图像上的缺陷区域，基本上将缺陷与电路板的其余部分分开。

    该函数首先使用OpenCV读取原始图像。在我们初始化与输入图像大小相同的蒙版之后。根据可视化标志，蒙版可以是3通道RGB蒙版（如彩色图像）或1通道灰度蒙版。最初，该掩码中的所有值都设置为零，这意味着没有缺陷。

  【3】缺陷标注解析。对于每个缺陷标注信息：识别缺陷类型，提取缺陷的多边形形状，该多边形被绘制到初始化的蒙版上。最后，该函数弥合了XML标注信息和适合训练的格式之间的差距。给定 PCB图像及其相应的XML 注释，它会生成一个分割掩模，突出显示有缺陷的区域。掩模可以是适合训练模型的数字格式，也可以是用于人工检查的视觉格式。

def process_mask(img_path, xml_path, visualize=False):
    img = cv2.imread(img_path)
    num_dim = 3 if visualize else 1
    mask = np.zeros((img.shape[0], img.shape[1], num_dim))


    # Read xml content from the file
    with open(xml_path, "r") as file:
      xml_content = file.read()


    data = xmltodict.parse(xml_content)


    # If has defect mask
    if "object" in data["annotation"]:
        objects = data["annotation"]["object"]


        # Single defects are annotated as a single dict, not a list
        if not isinstance(objects, List):
          objects = [objects]


        for obj in objects:
          defect_type = obj["name"]
          polygon = obj["polygon"]
          poly_keys = list(polygon.keys())


          # Get number of (x, y) pairs - polygon coords
          poly_keys = [int(k[1:]) for k in poly_keys]
          num_poly_points = max(poly_keys)


          # Parse ordered polygon coordinates
          poly_coords = []
          for i in range(1, num_poly_points+1):
              poly_coords.append([
                  int(float(polygon[f"x{i}"])),
                  int(float(polygon[f"y{i}"]))
              ])
          poly_coords = np.array(poly_coords)


          # Draw defect segment on mask
          fill_color = color_map[defect_type] if visualize else id_cat_map[defect_type]
          mask = cv2.fillPoly(mask, pts=[poly_coords], color=fill_color)


    #Optional
    if visualize:
        cv2.imwrite("output.jpg", mask)
        mask = Image.open("output.jpg")


    return mask

  【4】探索性数据分析。在训练模型之前，最好先了解数据。在这里，我们检查缺陷类型的分布并在样本图像上可视化缺陷。

    缺陷热力图显示了常见的缺陷位置，箱线图显示了缺陷尺寸的分布。

    该函数旨在通过读取边界框详细信息来计算 XML 注释中存在的每个缺陷的大小。

def get_defect_sizes(xml_paths):
    sizes = []
    for xml_path in xml_paths:
        with open(xml_path) as f:
            data = xmltodict.parse(f.read())
            objects = []
            if 'object' in data['annotation']:
                objects = data['annotation']['object']
            if not isinstance(objects, list):
                objects = [objects]
            for obj in objects:
                bndbox = obj['bndbox']
                width = int(bndbox['xmax']) - int(bndbox['xmin'])
                height = int(bndbox['ymax']) - int(bndbox['ymin'])
                sizes.append(width * height)
    return sizes

    最后，群图重点关注缺陷尺寸在整个数据集中的分布和扩散。

  【5】数据增强。该albumentations库用于人为扩展训练数据集，有助于提高模型的泛化能力。唯一指定的增强是水平翻转，它将以 50% 的概率水平翻转图像。

transform = A.Compose([
    A.HorizontalFlip(p=0.5)
])

 【6】图像预处理。将图像及其掩模预处理为适合Segformer模型的格式。

preprocessor = SegformerImageProcessor()

    我们将定义一个继承自PyTorchDataset类的自定义数据集类。这个自定义数据集允许我加载和预处理 PCB 图像及其相应的缺陷掩模。

    使用 OpenCV 加载图像。使用前面讨论的函数生成缺陷掩模process_mask。使用之前初始化的图像预处理图像及其掩模SegformerImageProcessor。此步骤将图像转换为张量格式，并确保它们具有适合 Segformer 模型的大小和标准化。返回预处理的图像和掩模张量。

class DefectSegmentationDataset(Dataset):
  def __init__(self, dataset, mode):
    self.image_paths = dataset["image_paths"]
    self.xml_paths = dataset["xml_paths"]


  def __len__(self):
    return len(self.image_paths)


  def __getitem__(self, idx):
    # Read image
    image = cv2.imread(self.image_paths[idx])


    # Get mask
    mask = process_mask(self.image_paths[idx], self.xml_paths[idx])
    mask = mask.squeeze()
    mask = Image.fromarray(mask.astype("uint8"), "L")


    # Return preprocessed inputs
    inputs = preprocessor(image, mask, ignore_index=None, return_tensors="pt")
    inputs["pixel_values"] = inputs["pixel_values"].squeeze(0)
    inputs["labels"] = inputs["labels"].squeeze(0)
    return inputs