记录UNet语义分割替换公开道路数据集训练的坑

最近在跑UNet模型，想要通过语义分割提取遥感影像里面的道路

代码是下载的一个B站up主的，讲解的很棒：科普：什么是语义分割_哔哩哔哩_bilibili

因为这个up主的代码没有公开道路的数据集，所以我下载了一个马萨诸塞州道路公开数据集，数据质量很高：道路和建筑物检测数据集 (toronto.edu)

因为我之前只做过目标检测，所以对语义分割不熟悉，过程中踩了一些坑：

1、np包的版本不同，导致np.int报错，修改成np.int_就行了；

2、FileNotFoundError报错，Up主的代码里图像是jpg格式，标签是png格式，要在代码里改成自己对应的格式：

3、特别注意：这里区分清楚图像分割结果的格式不是训练的图像的格式，是每次计算mIOU时都会生成的miou_out_path=".temp_miou_out"文件夹下的png格式图像，所以不要改成自己的训练图像格式了，要不每次运行到第5轮epoch计算mIOU的时候，都会报错；

4、"num_classes": 2, # 道路影像分类数量 num_classes是所需要区分的类的个数+1，这里我想要分一个类：道路，正常应该是"num_classes": 2,

如果训练的时候，设置"num_classes":1就会导致accuracy=0，mIOU=0，这肯定不正常，跑出来的效果就是这样的；

并且在预测时，"num_classes":1也会报错：

bash 复制代码

# ---------------------------------------------------#
#   进行图片的resize
# ---------------------------------------------------#
pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation=cv2.INTER_LINEAR)
print(pr.shape)
# ---------------------------------------------------#
#   取出每一个像素点的种类
# ---------------------------------------------------#
pr = pr.argmax(axis=-1)
print(pr.shape)

这里会报错：ValueError: cannot reshape array of size 4500 into shape (1500,1500,newaxis)，就是因为num_classes不对，所以在resize时shape不对，需要改正num_classes；

5、accuracy=100，mIOU=100，出现这种情况肯定也是不正常的，

原因可能是数据集的标签有问题，因为数据集的标签不对，模型识别不出来，或者是数据集没有标签；

6、如果需要训练自己的数据集，可以分为两种情况：

a、没有标签的数据集，按照视频里面的数据集标注教程，首先利用labelme标注图片生成json，然后转换成VOC格式的数据集；

b、有标签的数据集：像我这样下载网上常见的道路数据集，一般都是对输入图片分两类，背景的像素点值为0，目标（道路）的像素点值为255，这个像素值可以在GIS软件查看（马萨诸塞州道路公开数据集的道路就是255，背景是0）；

这样的数据集训练可以跑起来，但是预测是没有效果的，而且很可能导致训练中accuracy=100，mIOU=100；需要把标签图片的背景的像素值改为0，目标的像素值改为1，这个up主也写了对应的转换代码：GitHub - bubbliiiing/segmentation-format-fix: 该仓库中放置了一些工具，用于调整语义分割算法的数据集格式，使其符合训练要求。而且里面还有修改图像后缀名的代码；

修改之后的标签图片在文件夹中是全黑的，但是加载到GIS软件里可以看到，道路的像素是1，背景的像素是0，这样就可以训练了；

7、在代码中别忘了把预测中的num_classes改成自己数据集对应的数量（背景和道路）；

8、不要盲目修改在训练中输入图片的大小，改的太大了之后会爆显存；（注意，输入图像的大小要是32的倍数）

shell 复制代码

# --------------------------------#
#   输入图片的大小
# --------------------------------#
"input_shape": [512, 512],
# "input_shape": [1504, 1504], 
# 因为道路的影像宽高是1500像素，改大了会占用太多GPU，爆显存

其他的注意事项，例如labelme版本等等，在up主的代码里都有提到，可以自己去下载看一下；

以上就是我关于UNet模型替换公开道路数据集训练中的笔记，如果有不对的地方，可以找我交流。