U-Net实战 -- 隧道裂缝的识别
2024-05-10 12:21:29
U-Net网络学习算法，识别墙面裂缝
准备工作：
- Win10
- Git
- miniconda
- Pytorch
步骤
1. 下载裂缝数据集
Github原地址下载
2. 转换 mat 为 jpg
#!/user/bin/python
# coding=utf-8import os
from os.path import isdir
import numpy as np
from PIL import Image
from scipy import ioif __name__ == '__main__':file_path = r"此处改为你自己下载后的 groundTruth 路径"png_img_dir = r"此处改为你要把转换后 jpg 保存路径"if not isdir(png_img_dir):os.makedirs(png_img_dir)image_path_lists = os.listdir(file_path)images_path = []for index in range(len(image_path_lists)):image_file = os.path.join(file_path, image_path_lists[index])images_path.append(image_file)image_mat = io.loadmat(image_file)segmentation_image = image_mat['groundTruth']['Segmentation'][0]segmentation_image_array = np.array(segmentation_image[0])image = Image.fromarray((segmentation_image_array - 1) * 255)png_image_path = os.path.join(png_img_dir, "%s.jpg" % image_path_lists[index][0:3])image.save(png_image_path)# plt.figure()# plt.imshow(image)# plt.pause(2)# plt.show()
3. U-Net 代码
#!/user/bin/python
# coding=utf-8
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import lr_scheduler, optimizer
import torchvision
import os, sys
import cv2 as cv
from torch.utils.data import DataLoader, samplerclass SegmentationDataset(object):def __init__(self, image_dir, mask_dir):self.images = []self.masks = []files = os.listdir(image_dir)sfiles = os.listdir(mask_dir)for i in range(len(sfiles)):img_file = os.path.join(image_dir, files[i])mask_file = os.path.join(mask_dir, sfiles[i])# print(img_file, mask_file)self.images.append(img_file)self.masks.append(mask_file)def __len__(self):return len(self.images)def num_of_samples(self):return len(self.images)def __getitem__(self, idx):if torch.is_tensor(idx):idx = idx.tolist()image_path = self.images[idx]mask_path = self.masks[idx]else:image_path = self.images[idx]mask_path = self.masks[idx]img = cv.imread(image_path, cv.IMREAD_GRAYSCALE) # BGR ordermask = cv.imread(mask_path, cv.IMREAD_GRAYSCALE)# 输入图像img = np.float32(img) / 255.0img = np.expand_dims(img, 0)# 目标标签0 ~ 1， 对于mask[mask <= 128] = 0mask[mask > 128] = 1mask = np.expand_dims(mask, 0)sample = {'image': torch.from_numpy(img), 'mask': torch.from_numpy(mask), }return sampleclass UNetModel(torch.nn.Module):def __init__(self, in_features=1, out_features=2, init_features=32):super(UNetModel, self).__init__()features = init_featuresself.encode_layer1 = torch.nn.Sequential(torch.nn.Conv2d(in_channels=in_features, out_channels=features, kernel_size=3, padding=1, stride=1),torch.nn.BatchNorm2d(num_features=features),torch.nn.ReLU(),torch.nn.Conv2d(in_channels=features, out_channels=features, kernel_size=3, padding=1, stride=1),torch.nn.BatchNorm2d(num_features=features),torch.nn.ReLU())self.pool1 = torch.nn.MaxPool2d(kernel_size=2, stride=2)self.encode_layer2 = torch.nn.Sequential(torch.nn.Conv2d(in_channels=features, out_channels=features * 2, kernel_size=3, padding=1, stride=1),torch.nn.BatchNorm2d(num_features=features * 2),torch.nn.ReLU(),torch.nn.Conv2d(in_channels=features * 2, out_channels=features * 2, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 2),torch.nn.ReLU())self.pool2 = torch.nn.MaxPool2d(kernel_size=2, stride=2)self.encode_layer3 = torch.nn.Sequential(torch.nn.Conv2d(in_channels=features * 2, out_channels=features * 4, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 4),torch.nn.ReLU(),torch.nn.Conv2d(in_channels=features * 4, out_channels=features * 4, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 4),torch.nn.ReLU())self.pool3 = torch.nn.MaxPool2d(kernel_size=2, stride=2)self.encode_layer4 = torch.nn.Sequential(torch.nn.Conv2d(in_channels=features * 4, out_channels=features * 8, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 8),torch.nn.ReLU(),torch.nn.Conv2d(in_channels=features * 8, out_channels=features * 8, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 8),torch.nn.ReLU(),)self.pool4 = torch.nn.MaxPool2d(kernel_size=2, stride=2)self.encode_decode_layer = torch.nn.Sequential(torch.nn.Conv2d(in_channels=features * 8, out_channels=features * 16, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 16),torch.nn.ReLU(),torch.nn.Conv2d(in_channels=features * 16, out_channels=features * 16, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 16),torch.nn.ReLU())self.upconv4 = torch.nn.ConvTranspose2d(features * 16, features * 8, kernel_size=2, stride=2)self.decode_layer4 = torch.nn.Sequential(torch.nn.Conv2d(in_channels=features * 16, out_channels=features * 8, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 8),torch.nn.ReLU(),torch.nn.Conv2d(in_channels=features * 8, out_channels=features * 8, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 8),torch.nn.ReLU(),)self.upconv3 = torch.nn.ConvTranspose2d(features * 8, features * 4, kernel_size=2, stride=2)self.decode_layer3 = torch.nn.Sequential(torch.nn.Conv2d(in_channels=features * 8, out_channels=features * 4, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 4),torch.nn.ReLU(),torch.nn.Conv2d(in_channels=features * 4, out_channels=features * 4, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 4),torch.nn.ReLU())self.upconv2 = torch.nn.ConvTranspose2d(features * 4, features * 2, kernel_size=2, stride=2)self.decode_layer2 = torch.nn.Sequential(torch.nn.Conv2d(in_channels=features * 4, out_channels=features * 2, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 2),torch.nn.ReLU(),torch.nn.Conv2d(in_channels=features * 2, out_channels=features * 2, kernel_size=3, padding=1,stride=1),torch.nn.BatchNorm2d(num_features=features * 2),torch.nn.ReLU())self.upconv1 = torch.nn.ConvTranspose2d(features * 2, features, kernel_size=2, stride=2)self.decode_layer1 = torch.nn.Sequential(torch.nn.Conv2d(in_channels=features * 2, out_channels=features, kernel_size=3, padding=1, stride=1),torch.nn.BatchNorm2d(num_features=features),torch.nn.ReLU(),torch.nn.Conv2d(in_channels=features, out_channels=features, kernel_size=3, padding=1, stride=1),torch.nn.BatchNorm2d(num_features=features),torch.nn.ReLU())self.out_layer = torch.nn.Sequential(torch.nn.Conv2d(in_channels=features, out_channels=out_features, kernel_size=1, padding=0, stride=1),)def forward(self, x):enc1 = self.encode_layer1(x)enc2 = self.encode_layer2(self.pool1(enc1))enc3 = self.encode_layer3(self.pool2(enc2))enc4 = self.encode_layer4(self.pool3(enc3))bottleneck = self.encode_decode_layer(self.pool4(enc4))dec4 = self.upconv4(bottleneck)dec4 = torch.cat((dec4, enc4), dim=1)dec4 = self.decode_layer4(dec4)dec3 = self.upconv3(dec4)dec3 = torch.cat((dec3, enc3), dim=1)dec3 = self.decode_layer3(dec3)dec2 = self.upconv2(dec3)dec2 = torch.cat((dec2, enc2), dim=1)dec2 = self.decode_layer2(dec2)dec1 = self.upconv1(dec2)dec1 = torch.cat((dec1, enc1), dim=1)dec1 = self.decode_layer1(dec1)out = self.out_layer(dec1)return outdef test(unet):model_dict = unet.load_state_dict(torch.load('unet_road_model.pt'))root_dir = 'test'fileNames = os.listdir(root_dir)for f in fileNames:image = cv.imread(os.path.join(root_dir, f), cv.IMREAD_GRAYSCALE)h, w = image.shapeimg = np.float32(image) / 255.0img = np.expand_dims(img, 0)x_input = torch.from_numpy(img).view(1, 1, h, w)probs = unet(x_input.cuda())m_label_out_ = probs.transpose(1, 3).transpose(1, 2).contiguous().view(-1, 2)grad, output = m_label_out_.data.max(dim=1)output[output > 0] = 255predic_ = output.view(h, w).cpu().detach().numpy()# print(predic_)# print(predic_.max())# print(predic_.min())# print(predic_)# print(predic_.shape)# cv.imshow("input", image)result = cv.resize(np.uint8(predic_), (w, h))cv.imshow("unet-segmentation-demo", result)cv.waitKey(0)cv.destroyAllWindows()if __name__ == '__main__':index = 0num_epochs = 50train_on_gpu = Trueunet = UNetModel().cuda()# model_dict = unet.load_state_dict(torch.load('unet_road_model-100.pt'))# 如果要测试结果，取消下两行注释# test(unet)# exit()image_dir = './image'mask_dir = './gt8'dataloader = SegmentationDataset(image_dir, mask_dir)optimizer = torch.optim.SGD(unet.parameters(), lr=0.01, momentum=0.9)train_loader = DataLoader(dataloader, batch_size=1, shuffle=False)for epoch in range(num_epochs):train_loss = 0.0for i_batch, sample_batched in enumerate(train_loader):images_batch, target_labels = \sample_batched['image'], sample_batched['mask']print(target_labels.min())print(target_labels.max())if train_on_gpu:images_batch, target_labels = images_batch.cuda(), target_labels.cuda()optimizer.zero_grad()# forward pass: compute predicted outputs by passing inputs to the modelm_label_out_ = unet(images_batch)# print(m_label_out_)# calculate the batch losstarget_labels = target_labels.contiguous().view(-1)m_label_out_ = m_label_out_.transpose(1, 3).transpose(1, 2).contiguous().view(-1, 2)target_labels = target_labels.long()loss = torch.nn.functional.cross_entropy(m_label_out_, target_labels)print(loss)# backward pass: compute gradient of the loss with respect to model parametersloss.backward()# perform a single optimization step (parameter update)optimizer.step()# update training losstrain_loss += loss.item()if index % 100 == 0:print('step: {} \tcurrent Loss: {:.6f} '.format(index, loss.item()))index += 1# test(unet)# 计算平均损失train_loss = train_loss / dataloader.num_of_samples()# 显示训练集与验证集的损失函数print('Epoch: {} \tTraining Loss: {:.6f} '.format(epoch, train_loss))# test(unet)# save modelunet.eval()torch.save(unet.state_dict(), 'unet_road_model.pt')4. 测试结果
把待测裂缝图片截图为 480x320 后，在代码中开启测试代码运行。
参考文章：
https://www.jianshu.com/p/574f87abfc74
https://blog.csdn.net/ruotianxia/article/details/109703007
https://mp.weixin.qq.com/s/xeUdW2l71RsHe1Zdzr5a7Q
U-Net实战 -- 隧道裂缝的识别相关推荐
- [深度应用]·实战掌握Dlib人脸识别开发教程
[深度应用]·实战掌握Dlib人脸识别开发教程 个人网站--> http://www.yansongsong.cn/ 项目GitHub地址--> https://github.com/xi ...
- NLP实战-中文命名实体识别
NLP实战-中文命名实体识别:https://zhuanlan.zhihu.com/p/61227299
- 我的Go+语言初体验——Go+语言构建神经网络实战手写数字识别
"我的Go+语言初体验" | 征文活动进行中- 我的Go+语言初体验--Go+语言构建神经网络实战手写数字识别 0. 前言 1. 神经网络相关概念 2. 构建神经网络实战手写数字识 ...
- 前馈神经网络与支持向量机实战 --- 手写数字识别
前馈神经网络与支持向量机实战 - 手写数字识别 文章目录 前馈神经网络与支持向量机实战 --- 手写数字识别 一.前馈神经网络介绍 二.支持向量机介绍 三.数据集说明 四.环境准备 五.实验要求 六. ...
- pytorch 预测手写体数字_深度学习之PyTorch实战（3）——实战手写数字识别
如果需要小编其他论文翻译,请移步小编的GitHub地址 传送门:请点击我 如果点击有误:https://github.com/LeBron-Jian/DeepLearningNote 上一节,我们已经 ...
- opencv图像处理—项目实战：答题卡识别判卷
哔站唐宇迪opencv课程--项目实战:答题卡识别判卷 [计算机视觉-OpenCV]唐宇迪博士教会了我大学四年没学会的OpenCV OpenCV计算机视觉实战全套课程(附带课程课件资料+课件笔记+源码 ...
- MATLAB路面裂缝检测识别算法仿真
MATLAB路面裂缝检测识别算法仿真 在道路养护中,路面的裂缝是最常见的问题之一.对于路面裂缝的自动化检测和识别已经成为了一个研究热点.本文提出的基于MATLAB的路面裂缝检测识别算法能够减少人力.时 ...
- 信息抽取实战：命名实体识别NER【ALBERT+Bi-LSTM模型 vs. ALBERT+Bi-LSTM+CRF模型】（附代码）
实战:命名实体识别NER 目录 实战:命名实体识别NER 一.命名实体识别(NER) 二.BERT的应用 NLP基本任务 查找相似词语 提取文本中的实体 问答中的实体对齐 三.ALBERT ALBER ...
- 深度学习入门系列6项目实战：声纳回声识别
大家好,我技术人Howzit,这是深度学习入门系列第六篇,欢迎大家一起交流! 深度学习入门系列1:多层感知器概述 深度学习入门系列2:用TensorFlow构建你的第一个神经网络 深度学习入门系列3: ...
最新文章
- IPv6扩展头部 (四) 分片头部 Fragment Header for IPv6
- 陆奇给工程师们的5个建议
- MTK 驱动开发（39 ）---待机流程
- C语言中 malloc函数介绍
- matlab 混沌工具箱,matlab混沌工具箱
- Linux报错:tar: Error Is Not Recoverable: Exiting Now
- ios 区分iphone ipod ipad的方法及获取设备名称。
- 【树莓派4B】安装Ubuntu Mate20.04+ROS Noetic+使用电脑自带的xrdp和VNC进行PC端远程控制
- php怎么统计流量,PHP学习笔记:php网络流量统计系统
- 偷偷修复漏洞 苹果要求研究员噤声
- 【分布式系列】分布式事务：TX-LCN分布式事务框架整合实践
- 附录I-1 体系结构设计报告
- 【手撕代码】HDB3编解码
- Mac OS X上下载安装和配置hadoop指南
- 下载forestplot一直错误
- 快手小店怎么引流？快手怎么做店铺引流?
- w10打游戏老是弹出计算机,w10电脑打游戏时总是打出字怎么办
- java装修装饰公司订单管理系统ssm框架#java毕业设计
- oracle体系结构基础
- 制作企业宣传片的意义。