【深度学习|PyTorch】基于 PyTorch 搭建 U-Net 深度学习语义分割模型——附代码及其解释！

【深度学习|PyTorch】基于 PyTorch 搭建 U-Net 深度学习语义分割模型——附代码及其解释！

【深度学习|PyTorch】基于 PyTorch 搭建 U-Net 深度学习语义分割模型——附代码及其解释！

论文地址： https://arxiv.org/abs/1505.04597
代码地址：https://github.com/jakeret/tf_unet

1.数据准备

语义分割任务的输入通常是图像以及对应的像素级标签（即每个像素的分类）。我们首先需要加载和预处理数据。

import torch
from torch.utils.data import DataLoader, Dataset
import torchvision.transforms as transforms
from PIL import Image
import os

class SegmentationDataset(Dataset):
    def __init__(self, image_dir, mask_dir, transform=None):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.transform = transform
        self.images = os.listdir(image_dir)

    def __len__(self):
        return len(self.images)

    def __getitem__(self, index):
        img_path = os.path.join(self.image_dir, self.images[index])
        mask_path = os.path.join(self.mask_dir, self.images[index])
        image = Image.open(img_path).convert("RGB")
        mask = Image.open(mask_path).convert("L")  # Assuming masks are grayscale

        if self.transform is not None:
            image = self.transform(image)
            mask = self.transform(mask)

        return image, mask

# 数据加载及预处理
image_dir = "path_to_images"
mask_dir = "path_to_masks"

transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])

dataset = SegmentationDataset