python裂缝检测_小黑的Python日记：Unet简单实现裂缝分割

大噶好，我系小黑喵

裂缝数据集

--project

main.py

--image

--train

--data

--groundTruth

--val

--data

--groundTruth

我手动将数据集做成这个格式，其中trian84张，val34张，都保存为了jpg图像。

Unet

Unet结构

需要修改dataset.py为自己的数据集，其他小小改动即可。

#dataset.py

import torch.utils.data as data

import PIL.Image as Image

import os

def make_dataset(rootdata,roottarget):#获取img和mask的地址

imgs = []

filename_data = [x for x in os.listdir(rootdata)]

for name in filename_data:

img = os.path.join(rootdata, name)

mask = os.path.join(roottarget, name)

imgs.append((img, mask))#作为元组返回

return imgs

class MyDataset(data.Dataset):

def __init__(self, rootdata, roottarget, transform=None, target_transform=None):

imgs = make_dataset(rootdata,roottarget)

self.imgs = imgs

self.transform = transform

self.target_transform = target_transform

def __getitem__(self, index):

x_path, y_path = self.imgs[index]

img_x = Image.open(x_path).convert('L')#读取并转换为二值图像

img_y = Image.open(y_path).convert('L')

if self.transform is not None:

img_x = self.transform(img_x)

if self.target_transform is not None:

img_y = self.target_transform(img_y)

return img_x, img_y

def __len__(self):

return len(self.imgs)

#main.py

import numpy as np

import torch

import argparse

from torch.utils.data import DataLoader

from torch import autograd, optim

from torchvision.transforms import transforms

from unet import Unet

from dataset import MyDataset

# 是否使用cuda

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

x_transforms = transforms.Compose([

transforms.ToTensor(),

transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

])

# mask只需要转换为tensor

y_transforms = transforms.ToTensor()

def train_model(model, criterion, optimizer, dataload, num_epochs=10):

for epoch in range(0,num_epochs):

print('Epoch {}/{}'.format(epoch, num_epochs - 1))

print('-' * 10)

dt_size = len(dataload.dataset)

epoch_loss = 0

step = 0

for x, y in dataload:

step += 1

inputs = x.to(device)

labels = y.to(device)

# zero the parameter gradients

optimizer.zero_grad()

# forward

outputs = model(inputs)

loss = criterion(outputs, labels)

loss.backward()

optimizer.step()

epoch_loss += loss.item()

print("%d/%d,train_loss:%0.3f" %

(step,

(dt_size - 1) // dataload.batch_size + 1, loss.item()))

print("epoch %d loss:%0.3f" % (epoch, epoch_loss))

torch.save(model.cpu().state_dict(), 'weights_%d.pth' % epoch)

return model

#训练模型

def train():

batch_size = 1

liver_dataset = MyDataset(

"image/train/data", "image/train/gt",transform=x_transforms, target_transform=y_transforms)

dataloaders = DataLoader(

liver_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

train_model(model, criterion, optimizer, dataloaders)

#显示模型的输出结果

def test():

liver_dataset = MyDataset(

"image/val/data", "image/val/gt", transform=x_transforms, target_transform=y_transforms)

dataloaders = DataLoader(liver_dataset, batch_size=1)

import matplotlib.pyplot as plt

plt.ion()

with torch.no_grad():

for x, _ in dataloaders:

y = model(x)

img_y = torch.squeeze(y).numpy()

plt.imshow(img_y)

plt.pause(0.01)

plt.show()

if __name__ == '__main__':

pretrained = False

model = Unet(1, 1).to(device)

if pretrained:

model.load_state_dict(torch.load('./weights_4.pth'))

criterion = torch.nn.BCELoss()

optimizer = optim.Adam(model.parameters())

train()

test()

unet.py不需要变动

结果

训练了10个epoch后：累加loss大概到3

前几张预测图片：

上为预测，下为groundTruth

对于100多张的数据集，这个效果还行。

也算是填了一个以前的坑。