The code uses guide b.
The whole training process is also based on the video of guide b.
Source address: https://github.com/bubbliiiiing/unet-pytorch
Blog address: https://blog.csdn.net/weixin_44791964/article /details/108866828
# 1. Prepare the dataset
1. Use the labelme software to label the data and getjson文件
Note: The image format is .jpg, and the bit depth is 24 bits, otherwise it cannot be marked.
2. The format of the data set.
SegmentationClass stores the annotated png images of the pictures in JPEGImages.
3. Make your own data set
and put the original image in datasets>beforebefore
, which contains the original jpg image and json annotation file
Use labelme to generate json annotation files
4. After running json_to_dataset.py
to complete the annotation of the picture, you can run the file in the root directory json_to_dataset.pyto convert the data into a fixed format.
Need to modify the class before running
json_to_dataset.pyAfter the operation is completed, saved 1.jpg and 1.pngthey are stored in JPEGImages and SegmentationClass respectively
Finally, you can copy the JPEGImags and SegmentationClass files generated in the dataset to the JPEGImags and SegmentationClass files in VOCdevkit.
5. Divide the training set and verification set Generate train.txt file and val.txt file after
running the filevoc_annotion.py
Two, training
Parameter settings in the train.py file
Cuda Whether to use cuda, no gpu can be set to False
The number of num_classes is the number of categories + 1, where 1 is the background
model_path pre-training weight
VOCdevkit_path='VOCdevkit' dataset path
2. Train your own dataset
python train.py
The files saved after training are in the logs folder of the root directory
3. Use the trained model to make predictions
in the unet.py file
Modify model_path and num_classes
model_path to point to the weight file num_classes under the logs folder
as the number of categories +1
Run the predict.py file to predict the picture.
You need to input the path of the picture you want to predict.
4. Calculation of Miou evaluation index
Modify model_path and num_classes
model_path to point to the weight file num_classes under the logs folder
as the number of categories + 1
and modify name_classes
run get_miou.py