大噶好，我系小黑喵
裂缝数据集
数据集地址：https://github.com/cuilimeng/CrackForest-dataset
结构：
--project
main.py
--image
--train
--data
--groundTruth
--val
--data
--groundTruth
我手动将数据集做成这个格式，其中trian84张，val34张，都保存为了jpg图像。
Unet
论文地址：http://www.arxiv.org/pdf/1505.04597.pdf
代码来源：https://github.com/JavisPeng/u_net_liver
上面代码中，作者将Unet运用于liver识别，和裂缝一样，都只有一个mask，因而我们可以直接使用上述代码。
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]) # 复活了，这里修改就没错误了
])
# 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多张的数据集，这个效果还行。
也算是填了一个以前的坑。