unet 学习笔记-2，测试代码部分添加sigmoid函数，使得推理结果和label更加接近

记录下学习过程，当前还不是最终版本，这些不是最终版本的代码也打算保留，早期的代码一般集成度比较低，更通俗易懂一些。

网络结果代码:unet.py

import torch
from torch import nn


class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, input):
        return self.conv(input)


class Unet(nn.Module):
    def __init__(self,in_ch,out_ch):
        super(Unet, self).__init__()

        self.conv1 = DoubleConv(in_ch, 64)
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = DoubleConv(64, 128)
        self.pool2 = nn.MaxPool2d(2)
        self.conv3 = DoubleConv(128, 256)
        self.pool3 = nn.MaxPool2d(2)
        self.conv4 = DoubleConv(256, 512)
        self.pool4 = nn.MaxPool2d(2)
        self.conv5 = DoubleConv(512, 1024)
        self.up6 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
        self.conv6 = DoubleConv(1024, 512)
        self.up7 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.conv7 = DoubleConv(512, 256)
        self.up8 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.conv8 = DoubleConv(256, 128)
        self.up9 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.conv9 = DoubleConv(128, 64)
        self.conv10 = nn.Conv2d(64, out_ch, 1)

    def forward(self, x):
        c1 = self.conv1(x)
        p1 = self.pool1(c1)
        c2 = self.conv2(p1)
        p2 = self.pool2(c2)
        c3 = self.conv3(p2)
        p3 = self.pool3(c3)
        c4 = self.conv4(p3)
        p4 = self.pool4(c4)
        c5 = self.conv5(p4)
        up_6 = self.up6(c5)
        merge6 = torch.cat([up_6, c4], dim=1)
        c6=self.conv6(merge6)
        up_7 = self.up7(c6)
        merge7 = torch.cat([up_7, c3], dim=1)
        c7 = self.conv7(merge7)
        up_8 = self.up8(c7)
        merge8 = torch.cat([up_8, c2], dim=1)
        c8 = self.conv8(merge8)
        up_9 = self.up9(c8)
        merge9 = torch.cat([up_9, c1], dim=1)
        c9 = self.conv9(merge9)
        c10 = self.conv10(c9)
        return c10

网络训练代码：train.py

# -*- coding: utf-8 -*-

import torch
from torchvision.transforms import transforms
from torch import nn, optim

from unet import Unet
import numpy as np

# from tqdm import tqdm

import os
import cv2
import matplotlib.pyplot as plt
import numpy as np

path = '/home/kitty/BiSeNet-master/data/tongue'

train_image_path = os.path.join(path, 'train')
train_label_path = os.path.join(path, 'train_labels')
test_image_path = os.path.join(path, 'test')
test_label_path = os.path.join(path, 'test_labels')

train_image = os.listdir(train_image_path)
train_label = os.listdir(train_label_path)
test_image = os.listdir(test_image_path)
test_label = os.listdir(test_label_path)


PATH = './unet_model.pt'

# 是否使用cuda
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

x_transforms = transforms.Compose([
    transforms.ToPILImage(),
    # transforms.Resize(512,512),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

# mask 只需要转换为tensor
y_transforms = transforms.Compose([
    transforms.ToPILImage(),
    # transforms.Resize(512,512),
    transforms.ToTensor(),
    ])



def train_model(model, criterion, optimizer, num_epochs=10):
    best_model = model
    min_loss = 1000

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs-1))
        
        epoch_loss = 0
        
        step = 0
        
        for i in range(len(train_image)):
            step += 1
            
            
            image = cv2.imread(train_image_path + '/' + train_image[i], cv2.IMREAD_COLOR)
            image1 = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image1 = cv2.resize(image1, (512, 512))
            label = cv2.imread(train_label_path + '/' + train_label[i], cv2.IMREAD_GRAYSCALE)
            # label1 = cv2.cvtColor(label, cv2.COLOR_BGR2RGB)
            label1 = cv2.resize(label, (512, 512))
            
            
            inputs = x_transforms(image1).unsqueeze(0).to(device)
            labels = y_transforms(label1).unsqueeze(0).to(device)
            # inputs = x.to(device)
            # labels = y.to(device)
            

            optimizer.zero_grad()

            outputs = model(inputs)
            
            loss = criterion(outputs, labels)
            
            loss.backward()
            optimizer.step()
            
            epoch_loss += loss.item()
            print("%d, train_loss:%0.3f" % (step, loss.item()))
        print("epoch %d loss:%0.3f" % (epoch, epoch_loss/step))
        
        if (epoch_loss/step) < min_loss:
            min_loss = (epoch_loss/step)
            best_model = model
    torch.save(best_model.state_dict(), PATH)
    return best_model

# 训练模型
def train():
    model = Unet(3,1).to(device)
    # batch_size = 1
    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.Adam(model.parameters())
    
    # train_dataset = TrainDataset("", "", transform=x_transforms, target_transform=y_transforms())
    # dataloaders = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    
    train_model(model, criterion, optimizer)
    
def test():
    model = Unet(3,1)
    model.load_state_dict(torch.load(PATH))
    
    with torch.no_grad():
        for i in range(1):
            image = cv2.imread(test_image_path + '/' + test_image[i], cv2.IMREAD_COLOR)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = cv2.resize(image, (512, 512))
            print('image:', image.shape)
            inputs = x_transforms(image).unsqueeze(0)
            print('inputs:', inputs, inputs.shape)
            y = model(inputs)
            y = y.squeeze(0)
            y = y.permute(1,2,0)
            # plt.imshow(np.array(y))
            
            cv2.imshow('imshow', y.numpy()*255)
            # cv2.waitKey(0)
            print(y)

            plt.show()

            return y

if __name__ == '__main__':
    print("开始训练")
    # train()
    print("训练完成，保存模型")
    print("-"*20)
    print("开始预测")
    y = test()
    y = y.numpy()
    tt = y > -5
    tt = tt.astype(np.uint8)*255

    # cv2.imshow('tt', tt)
    # cv2.waitKey(0)


            

测试代码：test.py

# -*- coding: utf-8 -*-

import os
import cv2
import matplotlib.pyplot as plt

import torch
import torchvision.transforms as transforms

path = '/home/kitty/BiSeNet-master/data/tongue'

train_image_path = os.path.join(path, 'train')
train_label_path = os.path.join(path, 'train_labels')

train_image = os.listdir(train_image_path)
train_label = os.listdir(train_label_path)

# print(train_image)
#
# print(train_label)

img = cv2.imread(train_image_path + '/' + train_image[0], cv2.IMREAD_COLOR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# plt.figure(0)
#
# plt.imshow(img)
# plt.show()

x_transforms = transforms.Compose([
    transforms.ToPILImage(),
    # transforms.Resize(512, 512),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

img_t = x_transforms(img)

from unet import Unet
import numpy as np

PATH = 'unet_model.pt'
# 是否使用cuda
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def test():
    model = Unet(3,1).to(device)
    model.load_state_dict(torch.load(PATH))

    print(model)
    
    with torch.no_grad():
        for i in range(len(train_image)):
            image = cv2.imread(train_image_path + '/' + train_image[i], cv2.IMREAD_COLOR)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = cv2.resize(image, (512, 512))
            print('image:', image.shape)
            # inputs = x_transforms(image)
            inputs = x_transforms(image).unsqueeze(0).to(device)
            print('inputs.shape:', inputs.shape)
            import time
            print(time.time())
            y = model(inputs)
            print(time.time())
            y = y.squeeze(0)
            y = y.permute(1, 2, 0)

            # cv2.imshow('imshow', y.numpy() * 255)
            # cv2.waitKey(0)
            # print(y, y.shape)
            print(y)
            print('min,max', y.min(), y.max())
            y = torch.sigmoid(y)
            print('min,max', y.min(), y.max())
            print('min,max', y.min()*255, y.max()*255)
            y = y.cpu().numpy()
            print(y)

            # tt = y > -5
            tt = y
            tt = (tt * 255).astype(np.uint8)

            print(tt, tt.max())
            cv2.imshow('tt', tt)
            cv2.waitKey(3000)
            return 0

            # plt.show()

if __name__ == '__main__':

    print("开始预测")
    test()