分割方向API

分割方向API
https://github.com/qubvel/segmentation_models.pytorch
模型方面不需要写，只需要调用就可以了

以之前的 unet脑肿瘤分割完整代码为例
https://blog.csdn.net/qq_45845375/article/details/135588237

train2.py

import torch as t
import torch.nn as nn
from tqdm import tqdm  #进度条
import segmentation_models_pytorch as smp

from dataset import *


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

train_data=BrainMRIdataset(train_img,train_label,train_transformer)
test_data=BrainMRIdataset(test_img,test_label,test_transformer)

dl_train=DataLoader(train_data,batch_size=4,shuffle=True)
dl_test=DataLoader(test_data,batch_size=4,shuffle=True)


model = smp.Unet(
    encoder_name="resnet34",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
    encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
    in_channels=3,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
    classes=2,                      # model output channels (number of classes in your dataset)
)

img,label=next(iter(dl_train))
model=model.to('cuda')
img=img.to('cuda')
pred=model(img)
label=label.to('cuda')
loss_fn=nn.CrossEntropyLoss()#交叉熵损失函数
loss_fn(pred,label)
optimizer=torch.optim.Adam(model.parameters(),lr=0.0001)
def train_epoch(epoch, model, trainloader, testloader):
    correct = 0
    total = 0
    running_loss = 0
    epoch_iou = [] #交并比

    net=model.train()
    for x, y in tqdm(testloader):
        x, y = x.to('cuda'), y.to('cuda')
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        with torch.no_grad():
            y_pred = torch.argmax(y_pred, dim=1)
            correct += (y_pred == y).sum().item()
            total += y.size(0)
            running_loss += loss.item()

            intersection = torch.logical_and(y, y_pred)
            union = torch.logical_or(y, y_pred)
            batch_iou = torch.sum(intersection) / torch.sum(union)
            epoch_iou.append(batch_iou.item())

    epoch_loss = running_loss / len(trainloader.dataset)
    epoch_acc = correct / (total * 256 * 256)

    test_correct = 0
    test_total = 0
    test_running_loss = 0
    epoch_test_iou = []

    t.save(net.state_dict(), './Results2/weights/unet_weight/{}.pth'.format(epoch))

    model.eval()
    with torch.no_grad():
        for x, y in tqdm(testloader):
            x, y = x.to('cuda'), y.to('cuda')
            y_pred = model(x)
            loss = loss_fn(y_pred, y)
            y_pred = torch.argmax(y_pred, dim=1)
            test_correct += (y_pred == y).sum().item()
            test_total += y.size(0)
            test_running_loss += loss.item()

            intersection = torch.logical_and(y, y_pred)#预测值和真实值之间的交集
            union = torch.logical_or(y, y_pred)#预测值和真实值之间的并集
            batch_iou = torch.sum(intersection) / torch.sum(union)
            epoch_test_iou.append(batch_iou.item())

    epoch_test_loss = test_running_loss / len(testloader.dataset)
    epoch_test_acc = test_correct / (test_total * 256 * 256)#预测正确的值除以总共的像素点

    print('epoch: ', epoch,
          'loss： ', round(epoch_loss, 3),
          'accuracy:', round(epoch_acc, 3),
          'IOU:', round(np.mean(epoch_iou), 3),
          'test_loss： ', round(epoch_test_loss, 3),
          'test_accuracy:', round(epoch_test_acc, 3),
          'test_iou:', round(np.mean(epoch_test_iou), 3)
          )

    return epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc


if __name__ == "__main__":
    epochs=5
    for epoch in range(epochs):
        train_epoch(epoch,
                    model,
                    dl_train,
                    dl_test)

test2.py

import torch as t
import torch.nn as nn
import segmentation_models_pytorch as smp
from dataset import *
import matplotlib.pyplot as plt

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

train_data=BrainMRIdataset(train_img,train_label,train_transformer)
test_data=BrainMRIdataset(test_img,test_label,test_transformer)

dl_train=DataLoader(train_data,batch_size=4,shuffle=True)
dl_test=DataLoader(test_data,batch_size=4,shuffle=True)

model = smp.Unet(
    encoder_name="resnet34",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
    encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
    in_channels=3,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
    classes=2,                      # model output channels (number of classes in your dataset)
)

img,label=next(iter(dl_train))
model=model.to('cuda')
img=img.to('cuda')
pred=model(img)
label=label.to('cuda')
loss_fn=nn.CrossEntropyLoss()
loss_fn(pred,label)
optimizer=torch.optim.Adam(model.parameters(),lr=0.0001)
def test():
    image, mask = next(iter(dl_test))
    image=image.to('cuda')
    net = model.eval()
    net.to(device)
    net.load_state_dict(t.load("./Results2/weights/unet_weight/4.pth"))
    pred_mask = model(image)
    pred_mask=pred_mask
    mask=torch.squeeze(mask)
    pred_mask=pred_mask.cpu()
    num=4
    plt.figure(figsize=(10, 10))
    for i in range(num):
        plt.subplot(num, 4, i*num+1)
        plt.imshow(image[i].permute(1,2,0).cpu().numpy())
        plt.subplot(num, 4, i*num+2)
        plt.imshow(mask[i].cpu().numpy(),cmap='gray')#标签
        plt.subplot(num, 4, i*num+3)
        plt.imshow(torch.argmax(pred_mask[i].permute(1,2,0), axis=-1).detach().numpy(),cmap='gray')#预测
    plt.show()


if __name__ == "__main__":
    test()