语义分割宝库之segmentation_models_pytorch

最新推荐文章于 2025-11-17 03:29:29 发布

原创

最新推荐文章于 2025-11-17 03:29:29 发布 · 949 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #python #计算机视觉

本文介绍了segmentation_models_pytorch库，这是一个基于PyTorch的图像分割工具，特别适合图像分割任务。文中展示了如何利用CamVid数据集训练一个多类别语义分割模型，包括数据加载、图像增强、预处理、模型创建（Unet++）、损失函数选择（DiceLoss和CrossEntropyLoss）以及训练过程。

一、简介

segmentation_models_pytorch是一个基于PyTorch的图像分割神经网络。这个集合由俄罗斯的程序员小哥Pavel Yakubovskiy一手打造，对于图像分割而言简直就是神器般的存在。打比赛上分神器。
github地址：https://github.com/qubvel/segmentation_models.pytorch

二、训练自己的多类别语义分割模型

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '5'

import os
import numpy as np
import cv2
import albumentations as albu
import torch
import segmentation_models_pytorch as smp
from torch.utils.data import DataLoader
from torch.utils.data import Dataset as BaseDataset

# ---------------------------------------------------------------
### 加载数据
# CamVid数据集中用于图像分割的所有标签类别
CLASSES = ['bkg', 'ywm', 'line']
class Dataset(BaseDataset):
    """CamVid数据集。进行图像读取，图像增强增强和图像预处理.
    Args:
        images_dir (str): 图像文件夹所在路径
        masks_dir (str): 图像分割的标签图像所在路径
        class_values (list): 用于图像分割的所有类别数
        augmentation (albumentations.Compose): 数据传输管道
        preprocessing (albumentations.Compose): 数据预处理
    """
 
    def __init__(
            self,
            images_dir,
            masks_dir,
            augmentation=None,
            preprocessing=None,
    ):
        self.ids = os.listdir(images_dir)
        s