掌桥专利:专业的专利平台
掌桥专利
首页

一种基于知识蒸馏的医学图像分割方法、系统、设备

文献发布时间:2024-04-18 19:58:53


一种基于知识蒸馏的医学图像分割方法、系统、设备

技术领域

本领域涉及智能医疗领域,具体涉及一种基于知识蒸馏的医学图像分割方法、系统、设备、计算机可存储介质。

背景技术

随着现代医学的不断发展,各种医疗影像设备应用于临床诊疗,CT图像、X光图像、超声图像、PET图像和MRI图像等均为医师诊断治疗发挥着重要的作用。医学图像分割的主要目的是从各类医学图像中自动分割出病灶区域,辅助医师更快、更准确地做出诊断,极大地提高了医师的诊断效率。因此,医学图像分割模型的分割效果就显得格外重要。与此同时,医学图像相对于自然图像而言,背景环境更为复杂,具体表现在结构边缘更为模糊、形状更为多变、纹理结构细节更为丰富,使得医学图像分割任务更加具有挑战性。随着深度学习在医学图像分割中的广泛应用,医学图像分割得到快速的发展。卷积神经网络的出现,尤其是全卷积神经网络的出现让深度学习模型具备了高速度、高精度和全自动的特点。这些方法通常需要大量的标注数据和计算资源来训练复杂的模型,如FCN、DeepLab、SegNeXt等。然而,这些模型在推理时需要较长的时间和较高的内存消耗,不适合用于实时的医学图像分割,并且由于医学图像的复杂性以及目标区域的多样性,导致医学图像在边缘区域分割不完全和图像上下文特征信息利用不充分的问题上受到限制。

为了解决推理时间过长,不具备实时性的问题,一种可行的方案是使用嵌入式AI处理器来加速宫腔镜图像目标检测。嵌入式AI处理器是一种专门设计用于执行人工智能任务的芯片,如NVIDIA Jetson Nano、Google Coral Edge TPU等。这些芯片具有低功耗、高性能、小尺寸等优点,可以在边缘设备上实现快速的图像处理和推理。然而,直接将深度学习模型部署到嵌入式AI处理器上仍然存在一些问题,如模型过大、计算量过高、精度损失等。

发明内容

针对上述问题,本文提出一种基于知识蒸馏的医学图像分割方法,利用更好的大模型训练轻量化的小模型,使得轻量化小模型获得更好的性能和精度,降低模型的参数量和计算量,具体方法如下:

获取医学图像;

将所述医学图像输至训练好的分割模型进行目标分割中得到目标图像;所述训练好的分割模型的训练步骤包括:

S1:获取医学图像数据集及标签;

S2:将所述医学图像数据集及标签分别输至教师模型和学生模型中进行训练得到教师模型参数和学生模型参数,学生模型通过学习教师模型的参数更新自身参数得到更新后的学生模型参数,其中,学生模型由编码模块、边界增强模块、解码模块组成,所述编码模块进行特征提取得到特征图,所述边界增强模块依据所述特征图依次进行特征处理得到边界损失参数和增强特征图,所述解码模块对特征图进行解码后与增强特征图融合作为学生模型的输出;

S3:判断更新后的学生模型参数是否达到预期值,若是则导出新学生模型作为分割模型,若否则回到S2。

进一步,从所述特征图得到所述边界损失参数的特征处理包括:提取所述编码模块的前N层特征图,N为大于等于1的自然数,将所述前N层特征图进行特征融合后做自注意力计算得到增强特征图,所述增强特征图结合真实边界共同计算出边界损失参数。

所述教师模型与所述学生模型通过计算语义相似度分别对每个图像生成一个像素到像素的像像关系矩阵、一个像素到区域的像区关系矩阵;并且通过计算不同图像之间的相似度生成跨图像关系矩阵。

进一步,所述跨图像关系矩阵是通过计算每一对图像关系矩阵的JS散度后转换为相似度值得到的,具体的公式为:

其中,

进一步,对所述跨图像关系矩阵进行知识蒸馏,知识蒸馏产生的损失参数为:

其中,L

进一步,老师模型的损失参数包括下列的一种或几种:所述知识蒸馏产生的损失参数、传统的知识蒸馏损失参数、分割损害参数,学生模型的损失参数包括下列的一种或几种:传统的知识蒸馏损失参数、分割损害参数、边界损失参数通过加权计算得到总损失参数,所述总损失参数对学生模型进行参数更新。

进一步,损失老师模型的参数与学生模型的参数通过加权计算得到总损失参数,上述学生模型通过所述总损失参数更新模型得到更新后的学生模型。

本申请的目的在于提供一种基于知识蒸馏的医学图像分割系统,包括:

数据获取单元:获取医学图像;

图像分割单元:将所述医学图像输至训练好的分割模型进行目标分割中得到目标图像;所述训练好的分割模型的训练步骤包括:

S1:获取医学图像数据集及标签;

S2:将所述医学图像数据集及标签分别输至教师模型和学生模型中进行训练,学生模型通过学习教师模型的参数更新自身参数得到新学生模型参数,其中,学生模型由编码模块、边界增强模块、解码模块组成,所述编码模块进行特征提取得到特征图,所述边界增强模块依据所述特征图依次进行特征处理得到边界损失参数和增强特征图,所述解码模块对特征图进行解码后与增强特征图融合作为学生模型的输出;

S3:判断新学生模型参数是否达到预期值,若是则导出新学生模型作为分割模型,若否则回到S2。

本申请的目的在于提供一种基于知识蒸馏的医学图像分割设备,包括:

存储器与处理器,所述存储器用于存储程序指令;所述处理器用于调用程序指令,当程序指令被执行时实现任意一项上述的基于知识蒸馏的医学图像分割方法。

本申请的目的在于提供一种计算机可读存储介质,其上存储有计算机程序,所述计算机程序被处理器执行时实现任意一项上述的基于知识蒸馏的医学图像分割方法。

本申请的优势:

1.本申请采用知识蒸馏技术让轻量化小模型学习大模型的参数,使得小模型达到更优的性能,能够满足实时处理医学图像的检测分割的需求。

2.本申请采用了边界增强模块用于图像分割中增强边界特征,提高模型在分割中的准确率。

3.本申请在进行知识蒸馏的学生模型在更新参数时综合了边界增强损失参数、传统的知识蒸馏损失参数、传统的分割损失参数、蒸馏跨图像关系矩阵产生的损失参数,通过计算像素与像素之间以及像素与区域之间的相似度,将全局像素相关性从教师转移到学生进行语义分割,从而在整幅图像之间传递结构化的像素与像素、像素与区域之间的关系,使得蒸馏的方法可以较好地应用于目标检测任务。将所得性能良好的轻量化模型在AI处理器上进行部署,可以加快推理速度,满足医学图像检测和诊断对实时性的要求。

4.本申请通过对医学图像进行实时预处理和特征提取,然后使用经过大模型焦点和全局知识蒸馏的网络模型对图像中的特定组织、器官和病变区域等关键目标进行定位分割,最后将检测结果输出到显示器上,本发明能够有效地提高医学图像分割的速度和准确度,满足医学图像分割的实时性需求,降低推理时间和计算资源的消耗。

附图说明

为了更清楚地说明本发明实施例中的技术方案,下面将对实施例描述中所需要使用的附图作简单地介绍,显而易见地,下面描述中的附图仅仅是本发明的一些实施例,对于本领域技术人员来讲,在不付出创造性劳动的前提下,还可以根据这些附图获取其他的附图。

图1为本发明实施例提供的一种基于知识蒸馏的医学图像分割方法流程示意图;

图2为本发明实施例提供的一种基于知识蒸馏的医学图像分割系统示意图;

图3为本发明实施例提供的一种基于知识蒸馏的医学图像分割设备示意图;

图4为本发明实施例提供的模型构建过程的流程图;

图5为本发明实施例提供的学生模型BE-UNet的网络结构图;

图6为本发明实施例提供的边界增强模块结构图。

具体实施方式

为了使本技术领域的人员更好地理解本发明方案,下面将结合本发明实施例中的附图,对本发明实施例中的技术方案进行清楚、完整地描述。

在本发明的说明书和权利要求书及上述附图中的描述的一些流程中,包含了按照特定顺序出现的多个操作,但是应该清楚了解,这些操作可以不按照其在本文中出现的顺序来执行或并行执行,操作的序号如S101、S102等,仅仅是用于区分开各个不同的操作,序号本身不代表任何的执行顺序。另外,这些流程可以包括更多或更少的操作,并且这些操作可以按顺序执行或并行执行。需要说明的是,本文中的“第一”、“第二”等描述,是用于区分不同的消息、设备、模块等,不代表先后顺序,也不限定“第一”和“第二”是不同的类型。

图1本发明实施例提供的一种基于知识蒸馏的医学图像分割方法示意图,具体包括:

S101:获取医学图像;

在一个实施例中,医学图像包括:CT图像、X光图像、超声图像、PET图像和MRI图像,这些图像是当前疾病诊断过程中重要的诊断依据,通过计算机视觉技术、人工智能手段能够提供有效的诊断辅助,比如三维重建、病灶图像分割、人体姿态异常检测,临床疾病分型分类。

S102:将所述医学图像输至训练好的分割模型进行目标分割中得到目标图像;所述训练好的分割模型的训练步骤包括:

S1:获取医学图像数据集及标签;

S2:将所述医学图像数据集及标签分别输至教师模型和学生模型中进行训练,学生模型通过学习教师模型的参数更新自身参数得到新学生模型参数,其中,学生模型由编码模块、边界增强模块、解码模块组成,所述编码模块进行特征提取得到特征图,所述边界增强模块依据所述特征图依次进行特征处理得到边界损失参数和增强特征图,所述解码模块对特征图进行解码后与增强特征图融合作为学生模型的输出;

S3:判断新学生模型参数是否达到预期值,若是则导出新学生模型作为分割模型,若否则回到S2。

在一个实施例中,医学图像分割主要处理的是医学领域所涉及到的各种图像的分割问题,比如常见的核磁共振(MRI)扫描图像。其主要任务是,从这些医学图像中分割出兴趣区域,比如特定器官部位、兴趣目标(如肿瘤)等。医学图像分割与日常生活中常见场景的分割任务不同,医学图像(如MRI图像)由于图像获取设备的影响,会出现对比度低、信噪比低、光强低等问题;且器官本身存在运动和形变(如心脏),个体之间也有差异。这些因素导致了医学图像分割的难度和其自身的算法设计特点。

在一个实施例中,知识蒸馏是一种模型压缩方法,是一种基于“教师-学生网络思想”的训练方式,将复杂且大的模型作为教师模型,学生模型结构较为简单,用教师模型来辅助学生模型的训练,教师模型学习能力强,可以将它学到的知识迁移给学习能力相对弱的学生模型,以此来增强学生模型的泛化能力。复杂笨重但是效果好的教师模型不上线,就单纯是个导师角色,真正部署上线进行预测任务的是灵活轻巧的学生小模型;知识蒸馏的过程分为2个阶段:

1原始模型训练:训练"Teacher模型",它的特点是模型相对复杂,也可以由多个分别训练的模型集成而成。对"Teacher模型"不作任何关于模型架构、参数量、是否集成方面的限制,唯一的要求就是,对于输入X,其都能输出Y,其中Y经过softmax的映射,输出值对应相应类别的概率值。

2精简模型训练:训练"Student模型",它是参数量较小、模型结构相对简单的单模型。同样的,对于输入X,其都能输出Y,Y经过softmax映射后同样能输出对应相应类别的概率值。

在一个实施例中,所述教师模型与所述学生模型采用下列的一种或几种:SegGPT、SAM、SEEM、U-Net、SegNet、PSPNet、DeepLab v3+,学生模型中嵌入边界增强模块。

在一个实施例中,增强边界是加强图片中物体的边界,突出边界;不改变物体内部数据量,图像边缘上的灰度变化平缓,而边缘两侧灰度变化比较快,图像的边缘一般是指在局部不连续的图像特征,在图像分割中,为进一步提高分割性能,我们引入一个新的损失函数,称为边界增强损失,在训练过程中明确地关注边界区域。

在一个实施例中,知识蒸馏的损失函数可以表示为L=αL

在一个实施例中,在实际训练分割网络任务过程中,损失函数的选择尤为重要。对于语义分割而言,极有可能存在着正负样本不均衡,或者说类别不平衡的问题,因此选择一个合适的损失函数对于模型收敛以及准确预测有着至关重要的作用,常用的语义分割损失函数包括交叉熵、带权交叉熵、Focal loss、dice loss、IOU loss、Tversky loss。

在一个实施例中,所述特征图得到所述边界损失参数的特征处理包括:提取所述编码模块的前N层特征图,N为大于等于1的自然数,将所述前N层特征图进行特征融合后做自注意力计算得到增强特征图,所述增强特征图结合真实边界共同计算出边界损失参数。

在一个实施例中,所述教师模型与所述学生模型通过计算语义相似度分别对每个图像生成一个像素到像素的像像关系矩阵、一个像素到区域的像区关系矩阵;并且通过计算不同图像之间的相似度生成跨图像关系矩阵。

在一个实施例中,所述跨图像关系矩阵是通过计算每一对图像关系矩阵的JS散度后转换为相似度值得到的,具体的公式为:

其中,

在一个实施例中,对所述跨图像关系矩阵进行知识蒸馏,知识蒸馏产生的损失参数为:

其中,L

在一个实施例中,老师模型的损失参数包括下列的一种或几种:所述知识蒸馏产生的损失参数、传统的知识蒸馏损失参数、分割损害参数,学生模型的损失参数包括下列的一种或几种:传统的知识蒸馏损失参数、分割损害参数、边界损失参数通过加权计算得到总损失参数,所述总损失参数对学生模型进行参数更新。

在一个实施例中,损失老师模型的参数与学生模型的参数通过加权计算得到总损失参数,上述学生模型通过所述总损失参数更新模型得到更新后的学生模型。

在一个实施例中,本发明的模型构建过程如图4所示,首先输入医学图像,对医学图像进行预处理,将处理后的医学图像分别输至教师模型和学生模型中,其次,计算两个模型的像素到像素相似度值、像素到区域相似度值、跨图像区域相似度值、传统的知识蒸馏损失参数、传统的分割损失参数,然后对相似度值进行蒸馏得到像素到像素、像素到区域、跨图像区域的关系矩阵损失参数,此外,学生模型还结合边界增强模块得到边界增强损失参数,将边界增强损失参数、传统的知识蒸馏损失参数、传统的分割损失参数、像素到像素、像素到区域、跨图像区域的关系矩阵损失参数通过计算得到总损失参数,学生模型根据总损失参数进行模型参数更新,判断训练参数是否达到预期值,若是则将训练好的学生模型部署到AI处理器中用于医疗分割,若否则回到图像预处理阶段重复后续步骤。

在一个实施例中,本发明的学生模型进行参数更新时包括:

使用传统的知识蒸馏损失参数和分割损失参数+边界增强损失参数;

使用传统的知识蒸馏损失参数和分割损失参数+像素到像素、像素到区域、跨图像区域的关系矩阵损失参数;

使用传统的知识蒸馏损失参数和分割损失参数+边界增强损失参数+像素到像素、像素到区域、跨图像区域的关系矩阵损失参数。

在一个具体实施例中,本发明选择SegGPT作为教师网络,BE-UNet模型作为学生网络,它们都可以进行医学图像目标分割任务,例如分割肿瘤、器官、血管等,将图像输入分别输入到教师网络和BE-UNet中,得到特征图。其中,BE-UNet是在传统的UNet基础上增加了边界增强模块(Boundary Enhancement Module),该模块将UNet网络中前两层网络得到的特征图进行融合并进行注意力计算,随后再将得到的特征图与真实边界计算得到边界分割损失,并将此特征图与BE-UNet最终输出的特征图融合共同参与最终分割图的生成,具体的BE-UNet网络结构如图5所示,学生模型包括输入层、编码模块、边界增强模块、解码模块、输出层;输入层将图像数据输至编码模块,编码模块编码后将特征图输至解码模块,由解码模块输出解码后的特征图,此外,编码模块的前两层输出的特征图输至边界增强模块输出增减特征图,解码后的特征图与增强特征图融合后从输出层输出图像;边界增强模块的结构图如图6所示,编码模块前两层分别收到的特征图经过一维卷积或一维卷积+上采样后进行特征图融合得到融合特征图,对融合特征图进行自注意力计算得到增强特征图,将增强特征与真实边界相结合计算得到增强边界损失参数,并且增强特征图参与最终分割,即与解码后的特征图融合。

在一个具体实施例中,本发明为每个图像生成一个像素到像素的关系矩阵和一个像素到区域的关系矩阵,分别表示每个像素与其他像素或区域的语义相似度。具体来说,对于每个图像,首先使用老师网络和学生网络分别提取其特征图,然后将特征图划分为若干个区域,每个区域包含一定数量的像素。接着,计算每个像素与其他像素或区域的余弦相似度,得到两个关系矩阵,具体公式为:

其中,

然后,计算不同图像之间的关系矩阵的相似度,得到一个跨图像关系矩阵,表示每个图像与其他图像的结构化语义关系。具体来说,对于每一对图像,首先将它们的关系矩阵展平为向量,然后计算它们之间的JS散度,得到一个距离值。接着,将距离值转换为相似度值,得到一个跨图像关系矩阵,具体矩阵公式为:

其中,

再后,利用跨图像关系矩阵作为知识蒸馏的目标,让学生网络模仿老师网络在不同图像之间的结构化语义关系,结合传统的知识蒸馏损失和分割损失计算最终的损失函数。

其中,传统的知识蒸馏损失和分割损失由以下公式给出:

其中,L

其中,L

最终的L

L

其中,α,β和γ是三个超参数,用来控制不同部分损失函数的权重。

最后,将总损失用于优化学生网络的参数,使其尽可能接近教师网络的输出,同时保持良好的分割性能。通过这种方式,学生网络可以从教师网络中学习到更多的结构化语义关系知识,从而提高其对医学图像目标分割的能力;重复上述步骤(从教师模型和学生模型获得图像数据到学生模型更新参数),直到训练次数达到预期值;使用学生网络对新的医学图像进行目标分割,并评估其准确性和效率。

在一个具体实施例中,本发明从教师网络和学生网络的特征图中计算像素到像素的关系矩阵和像素到区域的关系矩阵,并使像素与像素,像素与区域的关系从教师网络传递到学生网络,使用跨图像关系损失来学习到教师网络中像素与像素和像素与区域之间的关系;使用传统的知识蒸馏损失和分割损失来学习教师网络的语义分割知识。

在一个具体实施例中,将训练得到的轻量最优模型部署到合适的AI处理器中,使用NNIE进行加速,利用AI处理器的强大加速能力加速异常检测过程,增强实时性。

在一个具体实施例中推理引擎可采用3588嵌入式AI处理器,经过大模型蒸馏优化过的模型可以方便地进行部署,在边缘设备上实现快速的图像处理和推理。

图2本发明实施例提供的一种基于知识蒸馏的医学图像分割系统示意图,具体包括:

数据获取单元:获取医学图像;

图像分割单元:将所述医学图像输至训练好的分割模型进行目标分割中得到目标图像;所述训练好的分割模型的训练步骤包括:

S1:获取医学图像数据集及标签;

S2:将所述医学图像数据集及标签分别输至教师模型和学生模型中进行训练,学生模型通过学习教师模型的参数更新自身参数得到新学生模型参数,其中,学生模型由编码模块、边界增强模块、解码模块组成,所述编码模块进行特征提取得到特征图,所述边界增强模块依据所述特征图依次进行特征处理得到边界损失参数和增强特征图,所述解码模块对特征图进行解码后与增强特征图融合作为学生模型的输出;

S3:判断新学生模型参数是否达到预期值,若是则导出新学生模型作为分割模型,若否则回到S2。

图3本发明实施例提供的一种基于知识蒸馏的医学图像分割设备示意图,具体包括:

存储器和处理器;所述存储器用于存储程序指令;所述处理器用于调用程序指令,当程序指令被执行任意一项上述的一种基于知识蒸馏的医学图像分割方法。

一种计算机可读存储介质,所述计算机可读存储介质存储计算机程序,所述计算机程序被处理器执行时任意一项上述的一种基于知识蒸馏的医学图像分割方法。

本验证实施例的验证结果表明,为适应症分配固有权重相对于默认设置来说可以改善本方法的性能。所属领域的技术人员可以清楚地了解到,为描述的方便和简洁,上述描述的系统,装置和单元的具体工作过程,可以参考前述方法实施例中的对应过程,在此不再赘述。本申请所提供的几个实施例中,应该理解到,所揭露的系统,装置和方法,可以通过其它的方式实现。例如,以上所描述的装置实施例仅仅是示意性的,例如,所述单元的划分,仅仅为一种逻辑功能划分,实际实现时可以有另外的划分方式,例如多个单元或组件可以结合或者可以集成到另一个系统,或一些特征可以忽略,或不执行。另一点,所显示或讨论的相互之间的耦合或直接耦合或通信连接可以是通过一些接口,装置或单元的间接耦合或通信连接,可以是电性,机械或其它的形式。所述作为分离部件说明的单元可以是或者也可以不是物理上分开的,作为单元显示的部件可以是或者也可以不是物理单元,即可以位于一个地方,或者也可以分布到多个网络单元上。可以根据实际的需要选择其中的部分或者全部单元来实现本实施例方案的目的。另外,在本发明各个实施例中的各功能单元可以集成在一个处理单元中,也可以是各个单元单独物理存在,也可以两个或两个以上单元集成在一个单元中。上述集成的单元既可以采用硬件的形式实现,也可以采用软件功能单元的形式实现。本领域普通技术人员可以理解上述实施例的各种方法中的全部或部分步骤是可以通过程序来指令相关的硬件来完成,该程序可以存储于一计算机可读存储介质中,存储介质可以包括:只读存储器(ROM,Read Only Memory)、随机存取存储器(RAM,Random AccessMemory)、磁盘或光盘等。

本领域普通技术人员可以理解实现上述实施例方法中的全部或部分步骤是可以通过程序来指令相关的硬件完成,所述的程序可以存储于一种计算机可读存储介质中,上述提到的介质存储可以是只读存储器,磁盘或光盘等。

以上对本发明所提供的一种计算机设备进行了详细介绍,对于本领域的一般技术人员,依据本发明实施例的思想,在具体实施方式及应用范围上均会有改变之处,综上所述,本说明书内容不应理解为对本发明的限制。

相关技术
  • 一种医学图像分割方法及装置、电子设备和存储介质
  • 一种基于人工智能的医学图像分割方法及系统
  • 一种基于知识蒸馏的BEV语义分割模型训练方法、系统、设备和介质
  • 一种基于知识蒸馏的花卉分类方法、系统、设备及介质
技术分类

06120116511906