devbox
小小林熬夜学编程
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

小地物小样本U-Net语义分割示例代码

作者：小小林熬夜学编程 | 2024-02-15 22:47:17

赞

踩

小地物小样本U-Net语义分割示例代码

仅需要准备数据，适当修改格式等（下面也有相应代码）

应该不挑环境，这里我用的是python3.8+cuda12.1+pytorch2.1.0，显卡是RTX3080

工作文件夹下建立data文件夹、model文件夹、utils文件夹、train.py、predict.py、其他一些代码（可选择性使用）

------------------------------------

data文件夹下建立test文件夹、train文件夹

test文件夹内直接放入要预测的图片，预测后的结果也在此文件夹中

train文件夹下建立image文件夹、label文件夹（都是图片）

要求/注意：图片均为PNG，512*512，三通道，image下位深可以不用管，label正常8通道，可把dataset.py中读取image后面加的函数复制过去也不再考虑位深，推荐label数据为二值化图像，读取也不会出错，image如果是从遥感影像tif格式转为png，且遥感影像有四通道（如近红波段），注意查看是否会有透明度的变化，因为png有RGBA 四通道，如果变了可以将原始的tif在Arcgis中图层右键导出数据，选择使用渲染器，强行RGB，格式选png导出。（但此方法未发现在Arcgis中批量导出的方法）

-----------------------------------

model文件夹内有三个.py文件：

__init__.py

unet_model.py

  1. """ Full assembly of the parts to form the complete network """
  2. """Refer https://github.com/milesial/Pytorch-UNet/blob/master/unet/unet_model.py"""
  3.  
  4. import torch.nn.functional as F
  5.  
  6. from .unet_parts import *
  7.  
  8.  
  9. class UNet(nn.Module):
  10.     def __init__(self, n_channels, n_classes, bilinear=True):
  11.         super(UNet, self).__init__()
  12.         self.n_channels = n_channels
  13.         self.n_classes = n_classes
  14.         self.bilinear = bilinear
  15.  
  16.         self.inc = DoubleConv(n_channels, 64)
  17.         self.down1 = Down(64, 128)
  18.         self.down2 = Down(128, 256)
  19.         self.down3 = Down(256, 512)
  20.         self.down4 = Down(512, 512)
  21.         self.up1 = Up(1024, 256, bilinear)
  22.         self.up2 = Up(512, 128, bilinear)
  23.         self.up3 = Up(256, 64, bilinear)
  24.         self.up4 = Up(128, 64, bilinear)
  25.         self.outc = OutConv(64, n_classes)
  26.  
  27.     def forward(self, x):
  28.         x1 = self.inc(x)
  29.         x2 = self.down1(x1)
  30.         x3 = self.down2(x2)
  31.         x4 = self.down3(x3)
  32.         x5 = self.down4(x4)
  33.         x = self.up1(x5, x4)
  34.         x = self.up2(x, x3)
  35.         x = self.up3(x, x2)
  36.         x = self.up4(x, x1)
  37.         logits = self.outc(x)
  38.         return logits
  39.  
  40. if __name__ == '__main__':
  41.     net = UNet(n_channels=3, n_classes=1)
  42.     print(net)

unet_parts.py

  1. """ Parts of the U-Net model """
  2. """https://github.com/milesial/Pytorch-UNet/blob/master/unet/unet_parts.py"""
  3.  
  4. import torch
  5. import torch.nn as nn
  6. import torch.nn.functional as F
  7.  
  8.  
  9. class DoubleConv(nn.Module):
  10.     """(convolution => [BN] => ReLU) * 2"""
  11.  
  12.     def __init__(self, in_channels, out_channels):
  13.         super().__init__()
  14.         self.double_conv = nn.Sequential(
  15.             nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
  16.             nn.BatchNorm2d(out_channels),
  17.             nn.ReLU(inplace=True),
  18.             nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
  19.             nn.BatchNorm2d(out_channels),
  20.             nn.ReLU(inplace=True)
  21.         )
  22.  
  23.     def forward(self, x):
  24.         return self.double_conv(x)
  25.  
  26.  
  27. class Down(nn.Module):
  28.     """Downscaling with maxpool then double conv"""
  29.  
  30.     def __init__(self, in_channels, out_channels):
  31.         super().__init__()
  32.         self.maxpool_conv = nn.Sequential(
  33.             nn.MaxPool2d(2),
  34.             DoubleConv(in_channels, out_channels)
  35.         )
  36.  
  37.     def forward(self, x):
  38.         return self.maxpool_conv(x)
  39.  
  40.  
  41. class Up(nn.Module):
  42.     """Upscaling then double conv"""
  43.  
  44.     def __init__(self, in_channels, out_channels, bilinear=True):
  45.         super().__init__()
  46.  
  47.         # if bilinear, use the normal convolutions to reduce the number of channels
  48.         if bilinear:
  49.             self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
  50.         else:
  51.             self.up = nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)
  52.  
  53.         self.conv = DoubleConv(in_channels, out_channels)
  54.  
  55.     def forward(self, x1, x2):
  56.         x1 = self.up(x1)
  57.         # input is CHW
  58.         diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
  59.         diffX = torch.tensor([x2.size()[3] - x1.size()[3]])
  60.  
  61.         x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
  62.                         diffY // 2, diffY - diffY // 2])
  63.  
  64.         x = torch.cat([x2, x1], dim=1)
  65.         return self.conv(x)
  66.  
  67.  
  68. class OutConv(nn.Module):
  69.     def __init__(self, in_channels, out_channels):
  70.         super(OutConv, self).__init__()
  71.         self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
  72.  
  73.     def forward(self, x):
  74.         return self.conv(x)

-----------------------------------

utils文件夹内有一个dataset.py文件：

  1. import torch
  2. import cv2
  3. import os
  4. import glob
  5. from torch.utils.data import Dataset
  6. import random
  7.  
  8.  
  9. class ISBI_Loader(Dataset):
  10.     def __init__(self, data_path):
  11.         # 初始化函数，读取所有data_path下的图片
  12.         self.data_path = data_path
  13.         self.imgs_path = glob.glob(os.path.join(data_path, 'image/*.png'))
  14.  
  15.     def augment(self, image, flipCode):
  16.         # 使用cv2.flip进行数据增强，filpCode为1水平翻转，0垂直翻转，-1水平+垂直翻转
  17.         flip = cv2.flip(image, flipCode)
  18.         return flip
  19.  
  20.     def __getitem__(self, index):
  21.         # 根据index读取图片
  22.         image_path = self.imgs_path[index]
  23.         # 根据image_path生成label_path
  24.         label_path = image_path.replace('image', 'label')
  25.         # 读取训练图片和标签图片
  26.         image = cv2.imread(image_path,flags=1)
  27.         label = cv2.imread(label_path)
  28.         # 将数据转为单通道的图片
  29.         image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  30.         label = cv2.cvtColor(label, cv2.COLOR_BGR2GRAY)
  31.         image = image.reshape(1, image.shape[0], image.shape[1])
  32.         label = label.reshape(1, label.shape[0], label.shape[1])
  33.         # 处理标签，将像素值为255的改为1
  34.         if label.max() > 1:
  35.             label = label / 255
  36.         # 随机进行数据增强，为2时不做处理
  37.         flipCode = random.choice([-1, 0, 1, 2])
  38.         if flipCode != 2:
  39.             image = self.augment(image, flipCode)
  40.             label = self.augment(label, flipCode)
  41.         return image, label
  42.  
  43.     def __len__(self):
  44.         # 返回训练集大小
  45.         return len(self.imgs_path)
  46.  
  47.  
  48. if __name__ == "__main__":
  49.     isbi_dataset = ISBI_Loader("data/train/")
  50.     print("数据个数：", len(isbi_dataset))
  51.     train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset,
  52.                                                batch_size=2,
  53.                                                shuffle=True)
  54.     for image, label in train_loader:
  55.         print(image.shape)

-----------------------------------

train.py

  1. from model.unet_model import UNet
  2. from utils.dataset import ISBI_Loader
  3. from torch import optim
  4. import torch.nn as nn
  5. import torch
  6.  
  7. def train_net(net, device, data_path, epochs=40, batch_size=1, lr=0.00001):
  8.     # 加载训练集
  9.     isbi_dataset = ISBI_Loader(data_path)
  10.     train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset,
  11.                                                batch_size=batch_size, 
  12.                                                shuffle=True)
  13.     # 定义RMSprop算法
  14.     optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.9)
  15.     # 定义Loss算法
  16.     criterion = nn.BCEWithLogitsLoss()
  17.     # best_loss统计，初始化为正无穷
  18.     best_loss = float('inf')
  19.     # 训练epochs次
  20.     for epoch in range(epochs):
  21.         # 训练模式
  22.         net.train()
  23.         # 按照batch_size开始训练
  24.         for image, label in train_loader:
  25.             optimizer.zero_grad()
  26.             # 将数据拷贝到device中
  27.             image = image.to(device=device, dtype=torch.float32)
  28.             label = label.to(device=device, dtype=torch.float32)
  29.             # 使用网络参数，输出预测结果
  30.             pred = net(image)
  31.             # 计算loss
  32.             loss = criterion(pred, label)
  33.             print('Loss/train', loss.item())
  34.             # 保存loss值最小的网络参数
  35.             if loss < best_loss:
  36.                 best_loss = loss
  37.                 torch.save(net.state_dict(), 'best_model.pth')
  38.             # 更新参数
  39.             loss.backward()
  40.             optimizer.step()
  41.  
  42. if __name__ == "__main__":
  43.     # 选择设备，有cuda用cuda，没有就用cpu
  44.     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  45.     # 加载网络，图片单通道1，分类为1。
  46.     net = UNet(n_channels=1, n_classes=1)
  47.     # 将网络拷贝到deivce中
  48.     net.to(device=device)
  49.     # 指定训练集地址，开始训练
  50.     data_path = './data/train/'
  51.     train_net(net, device, data_path)

-----------------------------------

predict.py

  1. import glob
  2. import numpy as np
  3. import torch
  4. import os
  5. import cv2
  6. from model.unet_model import UNet
  7.  
  8. if __name__ == "__main__":
  9.     # 选择设备，有cuda用cuda，没有就用cpu
  10.     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  11.     # 加载网络，图片单通道，分类为1。
  12.     net = UNet(n_channels=1, n_classes=1)
  13.     # 将网络拷贝到deivce中
  14.     net.to(device=device)
  15.     # 加载模型参数
  16.     net.load_state_dict(torch.load('best_model.pth', map_location=device))
  17.     # 测试模式
  18.     net.eval()
  19.     # 读取所有图片路径
  20.     tests_path = glob.glob('data/test/*.png')
  21.     # 遍历素有图片
  22.     for test_path in tests_path:
  23.         # 保存结果地址
  24.         save_res_path = test_path.split('.')[0] + '_res.png'
  25.         # 读取图片
  26.         img = cv2.imread(test_path)
  27.         # 转为灰度图
  28.         img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
  29.         # 转为batch为1，通道为1，大小为512*512的数组
  30.         img = img.reshape(1, 1, img.shape[0], img.shape[1])
  31.         # 转为tensor
  32.         img_tensor = torch.from_numpy(img)
  33.         # 将tensor拷贝到device中，只用cpu就是拷贝到cpu中，用cuda就是拷贝到cuda中。
  34.         img_tensor = img_tensor.to(device=device, dtype=torch.float32)
  35.         # 预测
  36.         pred = net(img_tensor)
  37.         # 提取结果
  38.         pred = np.array(pred.data.cpu()[0])[0]
  39.         # 处理结果
  40.         pred[pred >= 0.5] = 255
  41.         pred[pred < 0.5] = 0
  42.         # 保存图片
  43.         cv2.imwrite(save_res_path, pred)

-----------------------------------

剩下是一些杂七杂八的比较有针对性功能的代码：

格式转换tif变为png，但这个不能解决有透明度的问题，解决办法见文章最上面！

  1. from PIL import Image
  2. import os
  3.  
  4. def tif_to_png(input_path, output_path):
  5.     for file in os.listdir(input_path):
  6.  
  7.         if file.endswith('.tif'):
  8.             with Image.open(os.path.join(input_path, file)) as im:
  9.                 im.save(os.path.join(output_path, file.replace('.tif', '.png')))
  10.  
  11. # 示例
  12. tif_to_png('D:/deeplearning/tif',
  13.            'D:/deeplearning/png')

格式转换jpg变为png

  1. import os
  2. from PIL import Image
  3.  
  4.  
  5. # 获取指定目录下的所有png图片
  6. def get_all_png_files(dir):
  7.     files_list = []
  8.     for root, dirs, files in os.walk(dir):
  9.         for file in files:
  10.             if os.path.splitext(file)[1] == '.jpg':
  11.                 files_list.append(os.path.join(root, file))
  12.     return files_list
  13.  
  14.  
  15. # 批量转换png图片为jpg格式并保存到新的文件夹
  16. def png2jpg(files_list, output_dir):
  17.     for file in files_list:
  18.         img = Image.open(file)
  19.         new_file = os.path.splitext(file)[0] + '.png'
  20.         output_file = os.path.join(output_dir, os.path.basename(new_file))
  21.         img.convert('RGB').save(output_file)
  22.  
  23.  
  24. if __name__ == '__main__':
  25.     dir = r'D:\deeplearning\test' # png图片目录
  26.     output_dir = r'D:\deeplearning\test'  # 新的文件夹路径
  27.     files_list = get_all_png_files(dir)
  28.     png2jpg(files_list, output_dir)

将文件夹内图像批量重命名：

  1. #coding=gbk
  2. import os
  3. import sys
  4. def rename():
  5.     path=input("请输入路径(例如D:\\\\picture)：")
  6.     name=input("请输入开头名:")
  7.     startNumber=input("请输入开始数:")
  8.     fileType=input("请输入后缀名（如 .jpg、.txt等等）:")
  9.     print("正在生成以"+name+startNumber+fileType+"迭代的文件名")
  10.     count=0
  11.     filelist=os.listdir(path)
  12.     for files in filelist:
  13.         Olddir=os.path.join(path,files)
  14.         if os.path.isdir(Olddir):
  15.             continue
  16.         Newdir=os.path.join(path,name+str(count+int(startNumber))+fileType)
  17.         os.rename(Olddir,Newdir)
  18.         count+=1
  19.     print("一共修改了"+str(count)+"个文件")
  20.  
  21. rename()

参考包括：

Pytorch深度学习实战教程（三）：UNet模型训练 (qq.com)

GitHub：Deep-Learning/Tutorial/lesson-2 at master · Jack-Cherish/Deep-Learning · GitHub

B站up主：Bubbliiiing

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/小小林熬夜学编程/article/detail/87800
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号