语义分割宝库之segmentation_models_pytorch

 一、简介

        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