視網膜血管分割代碼(Pytorch實現)


創建日期: 2021-12-24 17:00:00

update log(2021.12.24):B站視頻刪除了,回放看了一下,講的不太行......2333,時間過得真快,轉眼就是2022年了啊

2021.08.01更新代碼講解視頻:

視網膜血管分割代碼分析(Pytorch實現)_嗶哩嗶哩_bilibili

簡介:

​ 本文主要分享我在做視網膜血管分割深度學習算法的過程中,整理出來的一套視網膜血管分割代碼。簡單介紹一下,我讀研期間的研究方向是計算機視覺(側重於語義分割),主要做一些醫學圖像分割算法和偏工程的機器視覺部署方面的工作,具體研究課題是基於深度學習的視網膜血管分割算法研究,選擇這個課題的原因也比較朴實------實驗室硬件計算資源有限。相比於深度學習算法創新,個人也更加偏向於算法落地應用方面,這個寒假還要惡補C++和計算機系統方面的知識。在這之前,把我的視網膜血管分割代碼分享給大家,項目地址:

https://github.com/lee-zq/VesselSeg-Pytorch

該項目是基於pytorch深度學習框架實現的視網膜血管分割代碼,包括數據預處理、模型訓練、模型測試以及可視化等功能,可以在此基礎上進一步研究視網膜血管分割算法。最近我把這套代碼進行了重構和簡化,在比較晦澀的地方也添加了注釋,力求任何一個入門者都能看懂。當然也可能存在bug和表述不清的地方,也希望大家能提issue指明問題,我會盡可能快速debug並更新代碼,不勝感激。(如果對視網膜血管分割比較熟悉,可以跳過下面的介紹,直接閱讀GitHub項目Readme)。Besides,Star⭐是我Git push的動力,后續我會搜集經典和SOTA的視網膜血管分割模型進行復現。

1. 相關知識

​ 基於深度學習做視網膜血管分割任務有一些相關的知識點可能需要說明一下,如果你剛入門建議先了解清楚。如果需要的話,我會單獨寫篇文章介紹。相關知識主要有以下幾點:

① 公開視網膜血管分割數據集。深度學習需要數據來驅動,發表文章需要和其他文獻的方法對比,所以建議使用公開視網膜血管分割數據集。視網膜血管分割的公開集最常用的有DRIVE、STARE和CHASE_DB1這三個,還有HRF等等不常用的數據集。當然,醫學圖像分割領域內私有數據還是比較認可的,也可以使用。

② 基於深度學習的圖像語義分割。這方面需要掌握基本的深度學習理論知識、Python和Pytorch深度學習框架,以及語義分割相關的知識點,如FCN和評價指標等等。

③ 視網膜血管分割研究現狀。這方面的工作現如今已完全深度學習化,深度學習方法的性能相比於血管跟蹤、模板匹配、形態學處理這類傳統方法要好很多,機器學習方法如今也逐漸被淘汰。

其他需要的也可以評論區指出。

2. 環境配置

​ 本項目代碼在Ubuntu20.04系統中使用Anaconda3創建的python虛擬環境運行,環境主要的包和版本如下:

# Name                    Version         
python                    3.7.9                    
pytorch                   1.7.0         
torchvision               0.8.0         
cudatoolkit               10.2.89       
cudnn                     7.6.5           
matplotlib                3.3.2              
numpy                     1.19.2        
opencv                    3.4.2         
pandas                    1.1.3        
pillow                    8.0.1         
scikit-learn              0.23.2          
scipy                     1.5.2           
tensorboardX              2.1        
tqdm                      4.54.1  

Pytorch1.0以上的版本兼容性很好,API基本沒變,所以可以先嘗試已存在的環境進運行。GPU為2080Ti,此外為了使訓練效率更高,訓練數據會一次性加載到內存,所以8G以下內存可能出出現內存溢出Error(具體取決於提取的patches數量)。

3. 使用方法

3.1 下載代碼

git clone https://github.com/lee-zq/VesselSeg-Pytorch.git

代碼結構以及功能注釋如下所示:

VesselSeg-Pytorch			# Source code		
    ├── config.py		 	# Configuration information
    ├── lib			            # Function library
    │   ├── common.py
    │   ├── dataset.py		        # Dataset class to load training data
    │   ├── extract_patches.py		# Extract training and test samples
    │   ├── help_functions.py		# 
    │   ├── __init__.py
    │   ├── logger.py 		        # To create log
    │   ├── losses
    │   ├── metrics.py		        # Evaluation metrics
    │   ├── pre_processing.py		# Data preprocessing
    ├── models		        # All models are created in this folder
    │   ├── denseunet.py
    │   ├── __init__.py
    │   ├── LadderNet.py
    │   ├── nn
    │   └── UNetFamily.py
    ├── prepare_dataset	        # Prepare the dataset (organize the image path of the dataset)
    │   ├── chasedb1.py
    │   ├── data_path_list		# image path of dataset
    │   ├── drive.py
    │   └── stare.py
    ├── test.py			            # Test file
    ├── tools			     # some tools
    │   ├── ablation_plot.py
    │   ├── ablation_plot_with_detail.py
    │   ├── merge_k-flod_plot.py
    │   └── visualization
    └── train.py			        # Train file

3.2 數據集准備

天翼雲盤鏈接下載我打包好的視網膜血管分割數據集,壓縮包包含DRIVE, CHASE_DB1和STARE三個數據集。或者你也可以從官網下載對應的數據集(DRIVESTARE和CHASE_DB1)。

將下載好的datasets.rar文件解壓到本地硬盤,得到的數據集目錄結構如下:

 datasets
    ├── CHASEDB1
    │   ├── 1st_label
    │   ├── 2nd_label
    │   ├── images
    │   └── mask
    ├── DRIVE
    │   ├── test
    │   └── training
    └── STARE
        ├── 1st_labels_ah
        ├── images
        ├── mask
        └── snd_label_vk

3.3 創建數據集路徑索引文件

項目根目錄下的"./prepare_dataset"目錄下有三個文件:drive.py,stare.py和chasedb1.py。分別將三個文件中的“data_root_path”參數賦值為上述3.2准備好的數據集的絕對路徑(例如: data_root_path="/home/lee/datasets")。然后分別運行:

python ./prepare_dataset/drive.py  
python ./prepare_dataset/stare.py
python ./prepare_dataset/chasedb1.py         

即可在"./prepare_dataset/data_path_list"目錄下對應的數據集文件夾中生成"train.txt"和"test.txt"文件,分別存儲了用於訓練和測試的數據路徑(每行依次存儲原圖,標簽和FOV路徑(用空格隔開))。

3.4 訓練模型

在根目錄下的"config.py"文件中修改超參數以及其他配置信息。特別要注意 “train_data_path_list"和"test_data_path_list"這兩個參數,分別指向3.3中創建的某一個數據集的"train.txt"和"text.txt"。 在"train.py"中構造創建好的模型(所有模型都在"./models"內手撕),例如指定UNet模型:

net = models.UNetFamily.U_Net(1,2).to(device)  # line 103 in train.py

修改完成后,在項目根目錄執行:

CUDA_VISIBLE_DEVICES=1 python train.py --save UNet_vessel_seg --batch_size 64

上述命令將在1號GPU上執行訓練程序,訓練結果保存在“ ./experiments/UNet_vessel_seg”文件夾中,batchsize取64,其余參數取config.py中的默認參數。

可以在config中配置培訓信息,也可以用命令行修改配置參數。訓練結果將保存到“ ./experiments”文件夾中的相應目錄(保存目錄的名稱用參數"--save"指定)。

此外,需要注意一點,config文件中有個“val_on_test”參數。當其為真時表示會在訓練的每個epoch結束后在測試集上進行性能評估,並選取"AUC of ROC"最高的模型保存為“best_model.pth”;當其為假時,會用驗證集性能評估結果(AUC of ROC)保存模型。當然保存最佳模型依據的指標可以自行修改,默認為AUC of ROC。

img

訓練過程示例(UNet)

3.5 測試評估

在“test.py”文件中構造對應的模型(同上),例如指定UNet模型:

net = models.UNetFamily.U_Net(1,2).to(device)

測試過程也需要"./config.py"中的相關參數,也可以在運行時通過命令行參數修改。

然后運行:

CUDA_VISIBLE_DEVICES=1 python test.py --save UNet_vessel_seg  

上述命令將訓練好的“./experiments /UNet_vessel_seg/best_model.pth”參數加載到相應的模型,並在測試集上進行性能測試,其測試性能指標結果保存在同一文件夾中的"performance.txt"中,同時會繪制相應的可視化結果。

img

測試過程示例(UNet)

img

DRIVE數據集預測結果性能指標(UNet)

可視化結果示例(LadderNet模型結果):

DRIVE數據集分割結果示例(依次為原圖、概率預測圖、二值預測圖和標簽圖):

img

img

img

img

DRIVE測試集ROC曲線

img

DRIVE測試集PR曲線

img

訓練過程驗證性能變化記錄

訓練過程的指標記錄也可通過Tensorboard日志文件查看:

img

img

img

img

以上。

  • 目前代碼整體上已經整理完畢,還有一些注釋和細微修改需要更新。有些我一個人考慮不周的地方還請指正。
  • 以上介紹部分或者相關內容有表述不詳盡的地方,還請提出來。有必要的話,我再寫一篇博客或者出一期視頻講解。

PS:該項目代碼參考了一些其他已發布代碼或設計方法,主要參考:


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM