深度学习处理图像分割的相关代码

以下是一个使用 PyTorch 实现语义分割的示例代码，具体实现了 U-Net 模型：

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms, datasets
from PIL import Image

# 定义 U-Net 模型
class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        self.down1 = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.down2 = nn.Sequential(
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True)
        )
        self.down3 = nn.Sequential(
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )
        self.down4 = nn.Sequential(
            nn.MaxPool2d(2),
            nn.Conv2d(256, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True)
        )
        self.bottom = nn.Sequential(
            nn.MaxPool2d(2),
            nn.Conv2d(512, 1024, 3, padding=1),
            nn.BatchNorm2d(1024),
            nn.ReLU(inplace=True),
            nn.Conv2d(1024, 1024, 3, padding=1),
            nn.BatchNorm2d(1024),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(1024, 512, 2, stride=2)
        )
        self.up1 = nn.Sequential(
            nn.Conv2d(1024, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(512, 256, 2, stride=2)
        )
        self.up2 = nn.Sequential(
            nn.Conv2d(512, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(256, 128, 2, stride=2)
        )
        self.up3 = nn.Sequential(
            nn.Conv2d(256, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(128, 64, 2, stride=2)
        )
        self.up4 = nn.Sequential(
            nn.Conv2d(128, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.out = nn.Conv2d(64, 2, 1)

    def forward(self, x):
        x1 = self.down1(x)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x4 = self.down4(x3)
        x5 = self.bottom(x4)
        x6 = self.up1(torch.cat([x5, x4], dim=1))
        x7 = self.up2(torch.cat([x6, x3], dim=1))
        x8 = self.up3(torch.cat([x7, x2], dim=1))
        x9 = self.up4(torch.cat([x8, x1], dim=1))
        x10 = self.out(x9)
        return x10

# 定义自定义数据集
class CustomDataset(Dataset):
    def __init__(self, root_dir):
        self.root_dir = root_dir
        self.transforms = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])

    def __getitem__(self, idx):
        img_path = self.root_dir + "/image_" + str(idx) + ".jpg"
        mask_path = self.root_dir + "/mask_" + str(idx) + ".jpg"

        img = Image.open(img_path).convert("RGB")
        img = self.transforms(img)

        mask = Image.open(mask_path).convert("L")
        mask = self.transforms(mask)

        return img, mask

    def __len__(self):
        return 100

# 定义训练函数
def train(model, dataloader, criterion, optimizer, device):
    model.train()

    running_loss = 0.0
    for i, (img, mask) in enumerate(dataloader):
        img = img.to(device)
        mask = mask.to(device)

        optimizer.zero_grad()

        output = model(img)
        loss = criterion(output, mask)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    return running_loss / len(dataloader)

# 加载数据集
dataset = CustomDataset("data")
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

# 定义模型、损失函数和优化器
model = UNet().cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(10):
    loss = train(model, dataloader, criterion, optimizer, "cuda")
    print("Epoch %d, Loss %.4f" % (epoch, loss))

在这个示例代码中，我们定义了一个 U-Net 模型，用于进行语义分割。我们还定义了一个自定义数据集 `CustomDataset`，用于加载训练数据。在训练过程中，我们将数据集分成大小为 4 的 batch 进行训练，并使用 Adam 优化器进行优化。