**************************************************** *
It's not easy to code words. Besides collecting, don't forget to give me a like!
**************************************************** *
---------Start
Please move to another blog post about the test part of the swinUnet network
Official code: https://github.com/HuCaoFighting/Swin-Unet
Purpose: Train Swin-Unet to segment lung regions
Official dataset location (may not be downloaded): https://www.kaggle.com/datasets/nikhilpandey360/chest-xray-masks-and-labels
CSDN free download dataset
Realize the effect:
input original image
output label
In this article, only 345 images from the entire dataset are used to complete the entire segmentation task!
1. Download the official code and unzip it
Unzipped folder:
2. Download the dataset and unzip it
We only need to use the following two folders: one for images and one for labels. There are 800 pictures in the original file, and only 704 labels. Some imgs have no labels. You need to pay attention when making npz files.
Here are the 345 images and corresponding masks used in this article
labels for segmentation tasks,
3. Generate .npz file
pycharm opens the project file, configures the python interpreter, and creates the data directory
In the data directory, train_npz is used to store the npz files used for training, and test_vol_h5 is used to store the npz files used for testing. This is the official name, and you can change the code less.
Convert images and labels to .npz files
Keep the original image and label in the same directory
Conversion code: (modify the path according to the location of your own data), if there are only two categories of background + target, this code can be used directly, if it is divided into three or more categories, the code should be based on Adjust your image data. After the adjustment, ensure that in the label array of the following code, the background is represented by 0 pixels, the target is represented by 1, 2, 3, 4... pixels respectively, and a pixel value represents a category. For example (0: background, 1: class 1, 2: class 2, 3: class 3...).
def npz():
#原图像路径
path = r'G:\dataset\Segmentation\LungSegmentation\npz\images\*.png'
#项目中存放训练所用的npz文件路径
path2 = r'G:\dataset\Unet\TransUnet-ori\data\Synapse\train_npz\\'
for i,img_path in enumerate(glob.glob(path)):
#读入图像
image = cv2.imread(img_path)
image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
#读入标签
label_path = img_path.replace('images','labels')
label = cv2.imread(label_path,flags=0)
#将非目标像素设置为0
label[label!=255]=0
#将目标像素设置为1
label[label==255]=1
#保存npz
np.savez(path2+str(i),image=image,label=label)
print('------------',i)
# 加载npz文件
# data = np.load(r'G:\dataset\Unet\Swin-Unet-ori\data\Synapse\train_npz\0.npz', allow_pickle=True)
# image, label = data['image'], data['label']
print('ok')
Generate training and testing image data as npz files in train_npz and test_vol_h5 respectively
Generate txt file corresponding to npz file
The content of the txt file is the name of the image data read in during model training and testing. Ignore the my_tools.py file.
The code to generate the txt file generates train.txt and test_vol.txt files according to the training and testing npz files, respectively.
def write_name():
#npz文件路径
files = glob.glob(r'C:\Users\22120\Desktop\Swin-Unet-main\data\Synapse\test_vol_h5\*.npz')
#txt文件路径
f = open(r'C:\Users\22120\Desktop\Swin-Unet-main\lists\lists_Synapse\test_vol.txt','w')
for i in files:
name = i.split('\\')[-1]
name = name[:-4]+'\n'
f.write(name)
4. Download pre-trained weights
csdn free download (recommended)
After the weights are downloaded, put them in the pretrained_ckpt folder of the project. The official only provides model weights with an input size of 224.
5. Modify part of the code
When your image data is single-channel, you can definitely train normally after modifying it according to what is written in the text. For input images with three channels and above, modify them according to the content written in the text. If there is still a problem, you can comment or private message me, and let’s solve it together.
5.1 Modify train.py
Modify the general parameters, the configuration file path, and note that num_classes is equal to the background + the number of predicted target categories. Because there are not many modifications, please forgive me for not putting the modified code, you can refer to the mark in the figure to modify it.
output_dir: The path where training logs and output weights are saved
root_path: The root directory where the dataset is stored
5.2 Modify dataset_synapse.py
The format of the npz file generated by myself is different from the official npz file format. It has been adjusted here, and it is completely the same after the adjustment.
5.3. Modify the trainer.py file
Set num_workers=0 in the DataLoader function in the trainer.py file
So far, all the codes have been modified, and the train.py file is executed. If the console has the following output, it will run successfully!
After training, the output folder in the project file stores the training output log and model weights.
