如何训练——盲道语义分割数据集，1400张左右，九个场景_附训练代码

计算机c9硕士算法工程师

1251人浏览 · 2024-11-05 08:13:11

计算机c9硕士算法工程师 · 2024-11-05 08:13:11 发布

数据准备

下载数据集
假设你已经从某个来源下载了盲道语义分割数据集，并将其解压到指定的目录中，例如 blind_path_dataset。
数据集结构
假设数据集的目录结构如下：

深色版本
blind_path_dataset/
├── images/
│ ├── 000001.jpg
│ ├── 000002.jpg
│ └── …
├── masks/
│ ├── 000001.png
│ ├── 000002.png
│ └── …
└── splits/
├── train.txt
├── val.txt
└── test.txt
images/ 目录包含输入图像。
masks/ 目录包含对应的标注图像，每个像素值表示一个类别。
splits/ 目录包含训练集、验证集和测试集的文件列表。
二、数据预处理

自定义数据集类
创建一个自定义的 PyTorch 数据集类来加载和预处理数据。

python
深色版本
import os
import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

自定义数据集类

class BlindPathDataset(Dataset):
def init(self, root_dir, split, transform=None):
self.root_dir = root_dir
self.split = split
self.transform = transform
self.image_paths = []
self.mask_paths = []

    with open(os.path.join(root_dir, 'splits', f'{split}.txt'), 'r') as f:
        lines = f.readlines()
        for line in lines:
            image_name = line.strip()
            self.image_paths.append(os.path.join(root_dir, 'images', f'{image_name}.jpg'))
            self.mask_paths.append(os.path.join(root_dir, 'masks', f'{image_name}.png'))

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

def __getitem__(self, idx):
    image = Image.open(self.image_paths[idx]).convert('RGB')
    mask = Image.open(self.mask_paths[idx])

    if self.transform:
        image = self.transform(image)
        mask = self.transform(mask)

    return image, mask

数据变换

transform = transforms.Compose([
transforms.Resize((256, 256)), # 调整图像大小
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 标准化
])

创建数据集实例

root_dir = ‘blind_path_dataset/’
train_dataset = BlindPathDataset(root_dir, ‘train’, transform=transform)
val_dataset = BlindPathDataset(root_dir, ‘val’, transform=transform)
test_dataset = BlindPathDataset(root_dir, ‘test’, transform=transform)

创建数据加载器

train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=4)
val_dataloader = DataLoader(val_dataset, batch_size=4, shuffle=False, num_workers=4)
test_dataloader = DataLoader(test_dataset, batch_size=4, shuffle=False, num_workers=4)

检查数据

for images, masks in train_dataloader:
print(images.shape, masks.shape)
break
三、模型选择
选择一个适合语义分割的模型，例如 U-Net。

python
深色版本
import torch
import torch.nn as nn
import torch.nn.functional as F

定义U-Net模型

class UNet(nn.Module):
def init(self, num_classes=9):
super(UNet, self).__4(self, num_classes=9)
self.num_classes = num_classes

    # 编码器
    self.enc1 = self.conv_block(3, 64)
    self.enc2 = self.conv_block(64, 128)
    self.enc3 = self.conv_block(128, 256)
    self.enc4 = self.conv_block(256, 512)
    self.pool = nn.MaxPool2d(2, 2)

    # 解码器
    self.upconv4 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
    self.dec4 = self.conv_block(512, 256)
    self.upconv3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
    self.dec3 = self.conv_block(256, 128)
    self.upconv2 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
    self.dec2 = self.conv_block(128, 64)
    self.upconv1 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
    self.dec1 = self.conv_block(64, 32)

    # 输出层
    self.out_conv = nn.Conv2d(32, num_classes, kernel_size=1)

def conv_block(self, in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True)
    )

def forward(self, x):
    enc1 = self.enc1(x)
    enc2 = self.enc2(self.pool(enc1))
    enc3 = self.enc3(self.pool(enc2))
    enc4 = self.enc4(self.pool(enc3))

    dec4 = self.upconv4(enc4)
    dec4 = torch.cat((dec4, enc3), dim=1)
    dec4 = self.dec4(dec4)
    dec3 = self.upconv3(dec4)
    dec3 = torch.cat((dec3, enc2), dim=1)
    dec3 = self.dec3(dec3)
    dec2 = self.upconv2(dec3)
    dec2 = torch.cat((dec2, enc1), dim=1)
    dec2 = self.dec2(dec2)
    dec1 = self.upconv1(dec2)
    dec1 = self.dec1(dec1)

    return self.out_conv(dec1)

实例化模型

model = UNet(num_classes=9)

将模型移到GPU

device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
model.to(device)
四、训练模型
定义损失函数、优化器，并进行模型训练。

python
深色版本

定义损失函数和优化器

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

训练模型

num_epochs = 10
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for images, masks in train_dataloader:
images, masks = images.to(device), masks.squeeze(1).long().to(device)

    optimizer.zero_grad()
    outputs = model(images)
    loss = criterion(outputs, masks)
    loss.backward()
    optimizer.step()
    
    running_loss += loss.item() * images.size(0)

epoch_loss = running_loss / len(train_dataset)
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}')

五、评估模型
使用验证集来评估模型的性能。

python
深色版本

评估模型

model.eval()
with torch.no_grad():
total_iou = 0.0
for images, masks in val_dataloader:
images, masks = images.to(device), masks.squeeze(1).long().to(device)
outputs = model(images)
preds = torch.argmax(outputs, dim=1)

    iou = self.compute_iou(preds, masks)
    total_iou += iou * images.size(0)

val_iou = total_iou / len(val_dataset)
print(f'Validation IoU: {val_iou:.4f}')

计算IoU

def compute_iou(preds, masks):
intersection = (preds & masks).float().sum((1, 2)) # 交集
union = (preds | masks).float().sum((1, 2)) # 并集
iou = (intersection + 1e-6) / (union + 1e-6) # 防止除零
return iou.mean()
六、保存模型
训练完成后，保存模型以供后续使用。

python
深色版本

保存模型

torch.save(model.state_dict(), ‘blind_path_segmentation_model.pth’)
七、完整代码
将上述所有步骤整合成一个完整的脚本：

python
深色版本
import os
import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

自定义数据集类

class BlindPathDataset(Dataset):
def init(self, root_dir, split, transform=None):
self.root_dir = root_dir
self.split = split
self.transform = transform
self.image_paths = []
self.mask_paths = []

    with open(os.path.join(root_dir, 'splits', f'{split}.txt'), 'r') as f:
        lines = f.readlines()
        for line in lines:
            image_name = line.strip()
            self.image_paths.append(os.path.join(root_dir, 'images', f'{image_name}.jpg'))
            self.mask_paths.append(os.path.join(root_dir, 'masks', f'{image_name}.png'))

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

def __getitem__(self, idx):
    image = Image.open(self.image_paths[idx]).convert('RGB')
    mask = Image.open(self.mask_paths[idx])

    if self.transform:
        image = self.transform(image)
        mask = self.transform(mask)

    return image, mask

数据变换

transform = transforms.Compose([
transforms.Resize((256, 256)), # 调整图像大小
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 标准化
])

创建数据集实例

创建数据加载器

定义U-Net模型

class UNet(nn.Module):
def init(self, num_classes=9):
super(UNet, self).init()
self.num_classes = num_classes

    # 编码器
    self.enc1 = self.conv_block(3, 64)
    self.enc2 = self.conv_block(64, 128)
    self.enc3 = self.conv_block(128, 256)
    self.enc4 = self.conv_block(256, 512)
    self.pool = nn.MaxPool2d(2, 2)

    # 解码器
    self.upconv4 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
    self.dec4 = self.conv_block(512, 256)
    self.upconv3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
    self.dec3 = self.conv_block(256, 128)
    self.upconv2 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
    self.dec2 = self.conv_block(128, 64)
    self.upconv1 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
    self.dec1 = self.conv_block(64, 32)

    # 输出层
    self.out_conv = nn.Conv2d(32, num_classes, kernel_size=1)

def conv_block(self, in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True)
    )

def forward(self, x):
    enc1 = self.enc1(x)
    enc2 = self.enc2(self.pool(enc1))
    enc3 = self.enc3(self.pool(enc2))
    enc4 = self.enc4(self.pool(enc3))

    dec4 = self.upconv4(enc4)
    dec4 = torch.cat((dec4, enc3), dim=1)
    dec4 = self.dec4(dec4)
    dec3 = self.upconv3(dec4)
    dec3 = torch.cat((dec3, enc2), dim=1)
    dec3 = self.dec3(dec3)
    dec2 = self.upconv2(dec3)
    dec2 = torch.cat((dec2, enc1), dim=1)
    dec2 = self.dec2(dec2)
    dec1 = self.upconv1(dec2)
    dec1 = self.dec1(dec1)

    return self.out_conv(dec1)

实例化模型

model = UNet(num_classes=9)

将模型移到GPU

device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
model.to(device)

定义损失函数和优化器

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

训练模型

num_epochs = 10
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for images, masks in train_dataloader:
images, masks = images.to(device), masks.squeeze(1).long().to(device)

    optimizer.zero_grad()
    outputs = model(images)
    loss = criterion(outputs, masks)
    loss.backward()
    optimizer.step()
    
    running_loss += loss.item() * images.size(0)

epoch_loss = running_loss / len(train_dataset)
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}')

评估模型

    iou = compute_iou(preds, masks)
    total_iou += iou * images.size(0)

val_iou = total_iou / len(val_dataset)
print(f'Validation IoU: {val_iou:.4f}')

计算IoU

保存模型

torch.save(model.state_dict(), ‘blind_path_segmentation_model.pth’)
八、运行脚本
将上述代码保存为一个 Python 脚本文件（例如 train_blind_path_segmentation.py），然后在命令行中运行：

bash
深色版本
python train_blind_path_segmentation.py
九、注意事项
数据集路径：确保数据集路径正确，特别是图片和标注文件的路径。
硬件资源：训练大型模型可能需要较多的计算资源，建议使用 GPU。
超参数调整：根据实际情况调整学习率、批量大小等超参数，以获得更好的训练效果。
模型选择：U-Net 是一个常用的语义分割模型，但你也可以尝试其他模型，如 DeepLabV3、FCN 等。
希望这些步骤能帮助你成功训练一个盲道语义分割模型。