Train your dataset with swinunet

Refer to the post ： https://blog.csdn.net/qq_37652891/article/details/123932772

Data set preparation

Multi category semantic segmentation of remote sensing images , The total is divided into 7 class （ Including background ）

image:

label_rgb

label( It's not all black here , The value of the category is 0,1,2,3,4,5,6), Later training uses such data

Data address
Baidu cloud ：https://pan.baidu.com/s/1zZHnZfBgVWxs6TJW4yjeeQ

Extraction code ：2022

SwinUNet Code address

Dataset processing

Data sets image and label, This data set should provide rgb Format labels and include 0,1,2,3,4,5,6 Label of the value ,SwinUNet The use is to include 0,1,2,3,4,5,6 Label image for ;

1. Data sets

Data sets are stored in SwinUNet The root directory ,image Middle is the original image ,label Middle is the label image （ common 7 class , Its tag value is 0,1,2,3,4,5,6,7）;
If you use other data sets , Pay attention to the value of the label . For example, if it is a second category . The label 0 or 255, It needs to be replaced with 0 or 1

—SwinUNet
---------configs
---------img_datas
---------------train
--------------------image
--------------------label
---------------test
--------------------image
--------------------label

2. stay SwinUnet Created in the root directory npz.py file , function npz.py file

import glob
import cv2
import numpy as np
import os

def npz(im, la, s):
    images_path = im
    labels_path = la
    path2 = s
    images = os.listdir(images_path)
    for s in images:
        image_path = os.path.join(images_path, s)
        label_path = os.path.join(labels_path, s)

        image = cv2.imread(image_path)
        image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
		#  Labels are converted from three channels to single channels 
        label = cv2.imread(label_path, flags=0)
        #  preservation npz file  
        np.savez(path2+s[:-4]+".npz",image=image,label=label)

npz('./img_datas/train/image/', './img_datas/train/label/', './data/Synapse/train_npz')

npz('./img_datas/test/image/', './img_datas/test/label/', './data/Synapse/test_vol_h5')

3. stay SwinUnet Created in the root directory txt.py file , function txt.py file

The goal is to generate ./list/list_Synapse/train.txt and ./list/list_Synapse/test_vol.txt file

import os
def write_name(np, tx):
    #npz File path 
    files = os.listdir(np)
    #txt File path 
    f = open(tx, 'w')
    for i in files:
        #name = i.split('\\')[-1]
        name = i[:-4]+'\n'
        f.write(name)
        
write_name('./data/Synapse/train_npz', './lists/lists_Synapse/train.txt')
write_name('./data/Synapse/test_vol_h5', './lists/lists_Synapse/test_vol.txt')

4. Download Pre Workout weights , Put it in SwinUnet In the catalog pretrained_ckpt Under the folder

link ：https://pan.baidu.com/s/1-hYwJRlr95Fv08e9AEARww
Extraction code ：2022

Modify the network

1. modify train.py file

Here's the important thing Number of categories , Others depend on the situation

2. modify ./datasets/dataset_synapse.py file

3. modify trainer.py file

I don't know why

4. Run code

This information can be passed in as a super parameter , If not , Then you can use default= To write the default value

If you set the default value , So run python train.py That's all right.