2024年YOLO还可以继续卷 | MedYOLO是怎么从YOLO家族中一步一步走过来的？

点击下方卡片，关注「集智书童」公众号

点击加入👉「集智书童」交流群

人工智能增强的医学影像器官、病变和其他结构的识别通常使用卷积神经网络（CNNs）设计，这些网络旨在对感兴趣的区域进行 Voxel 级精确分割。然而，用于训练这些CNN的标签需要花费大量时间生成，并且需要专家注意以确保质量。对于不需要 Voxel 级精度的任务，目标检测模型提供了一种可行的替代方案，可以减少标注的工作量。尽管这是一种潜在的应用，但目前尚无通用的3D医学影像检测框架可供选择。作者报告了MedYOLO，这是一种使用YOLO家族模型的一击检测方法设计的3D目标检测框架，适用于医学影像。

作者在四个不同的数据集上测试了这种模型：BRaTS，LIDC，一个腹部器官计算机断层扫描（CT）数据集和一个心电门控心脏CT数据集。作者发现，即使没有超参数调整，作者的模型在常见的中大尺寸结构（如心脏、肝脏和胰腺）上的性能也非常高。然而，模型在处理非常小或罕见的结构时存在困难。

Background

在3D医学影像中进行物体定位的标准方法是使用分割模型对感兴趣的目标进行 Voxel 到 Voxel 的标注。虽然这种方法使模型具有很高的准确性，但也存在一些缺点。为医学影像生成 Voxel 级准确的标注是一个耗时的过程，通常需要多个专家来验证标签的质量。由于标注者之间的变异性，器官或病变的医学术准确的分割可能会出现结构边界不确定的问题，这可能会导致附近组织中包含无关信息或排除相关信息。即使有高质量的标签，分割模型在准确标记目标结构边界时可能会遇到困难，通常需要后处理来填充缺失的内部体积并消除伪预测目标。总之，这使得分割模型的训练成本过高，同时可能会限制下游诊断或分类模型的预测能力。

目标检测模型提供了一种避免这些问题的替代方法，这些任务在 Voxel 特定精度不是必需的。然而，对于3D医学影像中的目标检测，相对的选择较少。专为照片设计的2D目标检测模型，如YOLO，可以提供按切片精确的边界框。但是，2D模型需要复杂的转换过程，既包括输入，以解决3D输入数据与模型不兼容的问题，也包括输出，以在下游3D任务中使用预测。此转换过程还丢弃了3D空间信息，这可能在检测复杂结构时有所帮助。

作者已经开发了一个用于医学影像的3D目标检测框架，以减轻这种选择稀少的问题。MedYOLO是基于Ultralytics YOLOv5检测模型开发的，对于中等和大型结构的检测具有很高的准确性。该框架与NIfTI成像具有原生兼容性。作者的结果表明，单击方法在具有基于 Patch 或滑动窗口方法的结构的检测中表现尤为出色，这些方法似乎在这些结构上表现不佳，如nnDetection[3]所使用的方法。

Methods

Data

作者在四种不同类型的医学影像上测试了MedYOLO。这些包括：

1. 来自BRaTS 2021任务1数据集[4, 5, 6, 7, 8]的FLAIR扫描（1000个训练扫描，251个验证），使用整个肿瘤分割掩膜生成边界框。

2. LIDC肺结节数据集（689个训练扫描，173个验证），使用两组不同的标签，一组将边界框放在单个结节周围，另一组使用包含每个扫描中每个结节的单个边界框。这些用于分别测试模型对非常小物体和扩散结构的检测。

1. 专有的腹部器官分割数据集（60个训练扫描，15个验证）。从该数据集中，为以下器官生成边界框：左侧肾脏、右侧肾脏、脾脏、胰腺、膈肌、膀胱、子宫、前列腺、主动脉、脊髓、胃和肝。对于这个数据集，还使用额外的预处理步骤生成旋转的训练示例。在轴视图上，使用每个示例（0，+-8和+-17度）的五个不同基础旋转角度将训练扫描旋转，每个示例还具有+-3度的额外随机扰动。对每个图像的分割掩膜应用相同的旋转，然后为旋转示例生成边界框标签。原始未旋转的示例也包含在训练集中，以减轻旋转算法产生的任何伪信号。总共，这为腹部数据集产生了360个训练示例。

2. 临床指示的随机选择的心电门控心脏计算机断层扫描（CT）数据集（648个训练扫描，163个验证），由梅奥诊所提供数据，并使用RILContour应用程序由作者标注，其中作者试图预测心脏和胸主动脉的边界框。每个数据集在患者 Level 上进行划分，以防止训练和验证集之间的数据泄漏。

Network Design

MedYOLO是基于YOLOv5框架开发的，专用于3D医学影像。它接受NIfTI文件作为输入，并包括CT和磁共振（MR）扫描的归一化功能。与YOLOv5一样，MedYOLO是一个基于 Anchor 点的检测模型，由一个连接到网络中较早层的中间层的卷积神经网络（CNN）的终端检测层创建。

与YOLOv5相比，MedYOLO的主要区别在于用其3D版本替换了2D神经网络层。MedYOLO的层数和宽度的配置可以通过YAML文件进行，其中包含小型、中型和大型版本。YOLOv5的一些关键库依赖，如OpenCV，与3D数据不兼容，需要对数据处理 Pipeline 和一些增强模式进行重大更改。MedYOLO代码库中的一部分与附加维度兼容，并保留了这部分代码以提高MedYOLO框架的可解释性。

与早期的YOLO版本一样，MedYOLO使用训练集标签进行k-means聚类来计算其 Anchor 框的先验。作者的初步测试使用肘方法确定了6个 Anchor 框，与YOLOv5使用的三个相比，被认为是适当数量。除测试过的每个数据集外，此设置保持不变，但这是一个可配置的超参数。

MedYOLO CNN需要类似于2D YOLO模型使用的立方输入，由于医学影像通常是各向同性的，作者使用双线性插值将输入数据转换为用户可配置大小的立方体。

为了在批处理大小与可用的GPU资源之间取得平衡，并避免由于输入大小不足而导致的错误，作者在大多数测试中选择了350像素边长的立方体。这使作者CNN底部的最终特征图大小为11x11x11。对于LIDC数据集，还测试了边长为512的立方体，这给出了最终特征图大小为16x16x16。请参阅表1以了解相应的GPU占用情况。

用于测试的数据集典型的NIfTI扫描形状通常在512x512x40和512x512x100之间，首先被加载到 Pipeline 中。该扫描被转换为PyTorch张量。然后，该张量从(X, Y, Z)转换为(Z, X, Y)，使得它具有形状40x512x512。然后，该张量被插值成一个形状为350x350x350的立方体。

接下来，根据下一节中的详细说明，应用增强。最后，应用适当的光学模态归一化。框架允许用户在需要时应用自己的归一化和预处理程序。

Training

MedYOLO的训练过程与YOLOv5非常相似。使用与YOLOv5从零开始训练时相同的超参数，唯一的例外是配置添加的增强程序的超参数。YOLOv5的一些增强程序需要删除，如颜色值变化和随机透视变换，因为它们适用于灰度数据或依赖于2D成像库。

使用了三种活增强方式：随机裁剪增强，其中将图像中的随机大小的块替换为随机噪声，随机平移增强和随机缩放增强。作者将MedYOLO的小版本在每一个数据集上训练了1000个Epoch，其中在200个Epoch后出现无改进时进行提前停止。

MedYOLO使用的是YOLOv5的复合损失函数的版本，该版本已经适应了3D体积。边界框损失组件比较预测和目标边界框的中心之间的交集与 union（IoU）以及距离。目标性组件，该组件训练模型通过将其预测的IoU与模型的置信度进行比较来评估其预测。最后，分类损失使用二进制交叉熵对预测的类别进行计算。存储库包含后两者（即目标性和分类损失）的焦损失选项，但本文中并未使用这些选项。

nnDetection框架是在相同的数据集上，使用与MedYOLO相同的训练和测试划分进行训练，以便进行比较。与MedYOLO一样，nnDetection也是一个用于3D NifTI图像的通用医学影像目标检测框架，但它使用的是滑动窗口目标检测方法，而不是MedYOLO的单次方法。

nnDetection框架也是自配置的，具有自动预处理、增强和5倍交叉验证。由于nnDetection的交叉验证是在预处理期间自动配置的，并且框架不跟踪患者ID，因此没有使用腹部CT数据集的旋转示例，以避免在不同折叠中出现来自单个患者的示例。

Results

在表2中，作者报告了在训练MedYOLO和nnDetection后获得的最佳验证指标。除非另有说明，否则报告的是MedYOLO的小版本（MedYOLO-S）。对于BRaTS数据集，在训练到第675个Epoch时出现了早停，对于ECG门控心脏数据集，在训练到第938个Epoch时出现了早停，而对于腹部数据集，小型模型训练了1000个Epoch。

作者还训练了MedYOLO的大版本（MedYOLO-L）在腹部数据集上，早停发生在第968个Epoch，在BRaTS数据集上，它未能收敛。作者对两个框架在IoU 0.5处的平均精度（mAP）进行了报告。

对于MedYOLO，作者报告了COCO指标的3D版本，该指标计算了0.5到0.95之间IoU的平均mAP，间隔为0.05，正如YOLOv5所报告的。nnDetection框架只报告了IoU间隔为0.1到0.9的mAP，因此对于作者，报告了一个稍微放松的指标，在0.5到0.9之间IoU的平均mAP，间隔为0.1。

得出的结果表明，MedYOLO在大尺寸、常见结构上的性能非常好，在大多数情况下超过了nnDetection的性能。特别是Gated Cardiac数据集，框架之间的行为有显著不同。尽管MedYOLO几乎完全捕获了目标体积，但nnDetection在高IoU处表现不佳。进一步调查发现，nnDetection框架的预测高度局部化到目标体积，但覆盖的目标体积很少。这导致了大量预测具有0.1到0.4之间的IoU，但对于这个数据集，没有预测具有IoU高于0.5的。

尽管在中等大小的和更大尺寸的结构上的结果良好，但MedYOLO在识别非常小结构和罕见类别的任务上遇到了困难。在训练期间，MedYOLO在测试的任何配置下都无法在LIDC数据集上展示任何学习进展。作者将这一现象归因于MedYOLO的CNN产生的最终特征图过于粗糙，限制了模型检测非常小结构的能力。

表1S显示了MedYOLO在Abdominal数据集上的性能比nnDetection更详细地划分。虽然MedYOLO在腹部通常优于nnDetection，但后者在识别子宫方面表现显著优于nnDetection，子宫在Abdominal训练示例中只出现17次，在验证示例中只出现4次。大型MedYOLO模型MedYOLO-L在提高Abdominal数据集上的性能方面提供了总体改进，但训练速度较慢。然而，它在子宫上的性能仍然显著低于nnDetection的性能。如训练策略中所讨论的，尽管MedYOLO-S与nnDetection的性能相似，但MedYOLO-L在BRaTS数据集上仍未能收敛。

Discussion

MedYOLO在作为3D医学影像的目标检测器方面显示出有前途的结果，能够使用相对简单的，一次性的检测方法达到良好的结果。与nnDetection相比，MedYOLO在中等尺寸的结构上表现出类似的效果，并在nnDetection使用的滑动窗口方法的大尺寸结构上表现出明显优于nnDetection的效果。

然而，在作者的分析中，它未能检测到非常小或弥散的结构。在不需要超参数调优的情况下，在各种结构上实现高mAP，表明MedYOLO对次优超参数选择的鲁棒性。框架能够快速、准确地找到中大型结构，使其成为机器学习 Pipeline 中用于将数据定位到相关体积的理想工具。

尽管在某些任务上表现出色，但当前实现具有很大的改进潜力。最大的改进途径之一是增加额外的增强程序。特别是，在BRaTS肿瘤检测上的表现，尽管MedYOLO-S取得了相当 modest 的结果并提前停止训练，而MedYOLO-L无法收敛，表明模型可能缺少相关信息。增强训练可以延长训练过程，以实现在这类和其他数据集上的更高性能。作者遇到的障碍是3D数据的相对缺乏高效且易于实现的增强程序，这导致许多标准的实时增强策略显著降低了训练速度。预生成增强训练数据可以避免这个问题，但由于医学影像的大内存占用，这会极大地增加存储要求。

另一个可能改进性能的修改是更换作者用于 Reshape 数据所使用的算法。三线性插值允许作者平滑地将3D输入数据转换为立方形状，但并未提高输入数据的信息。更复杂的插值方法，如超分辨率，可以提供额外的细节并增加在 Reshape 过程中创建的切片的价值。

将输入数据转换为立方体积的需求可能是作者 Pipeline 中最大的弱点。增加输入切片数量，通常是一个数量级，会大大增加作者模型的计算资源需求。此外，在平衡批处理内存约束和批处理统计学准确性的同时，作者需要使用较小的立方体并减少轴向分辨率。在多次批处理之间累积梯度可能有助于解决这个问题。但作者没有尝试这样做，因为在大多数测试数据集上，作者没有发现需要它来提高性能。尽管这对于LIDC数据集可能是有用的，但在即使给模型提供图像的全部轴向分辨率的情况下，模型在训练过程中仍然没有取得任何进展，这表明这种情况的可能性很小。

与医学影像数据集有关的一个相关问题，这些数据集通常包含具有不同切片数量的图像。将这些数据集重新缩放到固定大小的立方体可能会以不可预测的方式相对于彼此扭曲输入图像。在推理时，具有异常少切片的扫描是模型主要无法准确预测边界框的，即使对于作者的验证指标在ECG门控心脏数据集上极高的数据集也是如此。

在开发早期，作者投入了大量的努力，将立方形状的要求从MedYOLO框架中移除，至少对于增加的深度维度，但这是有必要的，原因如下。首先，MedYOLO和YOLO架构中的深度和池化层数量对每个维度都提出了最小分辨率要求。大多数讨论的数据集选择的分辨率，350x350x350，接近这个最小分辨率。

更难以解决的问题是，YOLOv5模型需要方形输入数据来简化基于 Anchor 点的预测。当将这些2D模型转换为3D时，这成为立方输入的必要条件。为了处理矩形输入，YOLOv5对输入数据进行重采样和填充以获得方形形状，这降低了重采样过程中的畸变。对于2D数据，这个过程是计算成本低的，但是医学影像的大内存占用和高度各向异性的形状使其成为将3D输入数据填充到立方体的实际困难，限制了作者只能使用重采样技术。

尽管存在这些局限性，MedYOLO证明了一个基于 Anchor 点的，一次性的方法可以在3D医学目标检测任务上实现高性能。然而，未来的框架可能使用这种类似于YOLO的2.5D范式来做得更好。这将使新的框架能够在不牺牲批处理大小或引入重采样畸变的情况下保持输入数据的原始分辨率。

与3D方法相比，2.5D方法的主要缺点是需要额外的标注工作来保持大型结构上的边界框精度，但与 Voxel 级分割相比，这仍然简化了标注工作。

Conclusion

作者报告了MedYOLO的开发，它是一个基于 Anchor 点的目标检测框架，专为使用3D医学影像而设计。MedYOLO在不需要超参数优化的情况下，可以匹配或超过nnDetection在定位中等到大尺寸结构方面的性能。Pipeline 未来可以改进的领域包括添加增强程序，以便在开发高效的3D算法时提高性能，以及将更复杂的插值方法集成到重采样输入数据以达到所需立方体形状中。

参考

[1].MedYOLO: A Medical Image Object Detection Framework .

扫码加入👉「集智书童」交流群

（备注：方向+学校/公司+昵称）

想要了解更多：

前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉（分割、检测）、多模态、目标跟踪、NerF」

行业技术方案👉「AI安防、AI医疗、AI自动驾驶」AI模型部署落地实战👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」

欢迎扫描上方二维码，加入「集智书童-知识星球」，日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑，期待交流！

免责声明凡本公众号注明“来源：XXX（非集智书童）”的作品，均转载自其它媒体，版权归原作者所有，如有侵权请联系我们删除，谢谢。
点击下方“阅读原文”，了解更多AI学习路上的「武功秘籍」