基于UNet网络的新冠肺炎分割系统

2024-03-167 min 分享
cover image
AI摘要
本文介绍了一种基于 UNet 网络的新冠肺炎分割系统。该系统利用 Pytorch 和 Pyqt6 开发,使用公开的 Covid-19 CT 图像数据集进行训练。系统可以根据患者的肺部 CT 图检测病灶面积并统计。为了训练模型,作者对原始数据集进行了二次处理。训练出的模型可以有效地分割病灶区域,但由于数据集较小,训练出可用模型需要调整训练参数。作者分享了处理好的数据集和训练好的模型,方便其他人使用。

前期专业方向实践的课题是基于深度学习的新冠肺炎分割系统,虽然没什么技术含量,但找数据集和训练模型确实是废了些功夫,写篇博客分享一下资源和心得体会

成品展示

我不知道其他学校有没有专业方向实践,反正我们学校这就相当于一个大课设,跟课设的难度和重视程度差不多,都不高,但学分很高,属于是大四给我们刷绩点的课吧。因为是深度学习,那肯定是 Pytorch 了,去 Github 上找了个 UNet 的开源训练框架,然后界面使用 pyqt6 写的,简简单单写了一下

功能大致就是根据患者的肺部CT图,检测病灶面积并统计

框架

我是直接 copy 的 milesial/Pytorch-UNet,进行了一些修改,增加了图形界面。Gmero/Pytorch-UNet-Covid19 (gitee.com)

请确保你已经搭建好了 pytorch 环境并安装了 CUDA,可以参考这个视频 PyTorch深度学习快速入门教程(绝对通俗易懂!)

通过 train.py 启动模型训练模块(具体参数参考 README),main.py 启动主程序(确保你已经准备好了模型和演示数据集),在 main.py 中找到 load_model 可以指定模型位置。如果你不想训练模型的话,可以下载我已经训练好的模型和数据集,往下看

数据集

数据集一开始卡了我很久,不管啥搜索引擎都是一些要你花钱买的。Github上找半天也是一些不带标签的CT图,没办法进行训练,碰运气问了一下 newBing,结果真找到了一个能用的 图像下载 - 2019新型冠状病毒信息库 (2019nCoVR) (big.ac.cn),官方网站,有大量的患者CT图,还有一份带标签的训练集(往下翻找到Experiment data files),训练集不大,但这是我唯一能找到的了

原始数据集的标签是通过灰度值分成背景、肺野、毛玻璃混浊 (GGO) 和实变 (CL)。由于这里面实变的数据特别少,学习效果很差,我将实变与毛玻璃浑浊进行了合并,就叫病灶吧。然后原始数据灰度值都是012这种,肉眼不可见,写了个脚本进行了二次处理。背景为纯黑色,肺部区域为灰色,病灶区域为白色

直链下载处理好的可用于训练的训练集和一份用于演示的患者CT图像 COVID19-PytorchUNet-dataset.zip

模型

一开始训练出的模型几乎是不可用的状态,可能是数据集太小了的原因,后来在高人指点的情况下对训练参数进行了调整 -e 100 -b 8 ,终于是训练出了能用的模型,我的4060ti跑冒烟了都

如果你们不想自己训练的话,可以下载我训练好的模型 checkpoint_epoch100.pth,放到根目录 checkpoints 文件夹,没有就新建

评论
正在加载评论组件...