使用qubvel/segmentation_models库进行图像分割的完整指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_01014/article/details/148508553

使用qubvel/segmentation_models库进行图像分割的完整指南

segmentation_models Segmentation models with pretrained backbones. Keras and TensorFlow Keras. 项目地址: https://gitcode.com/gh_mirrors/se/segmentation_models

项目概述

qubvel/segmentation_models是一个基于Keras/TensorFlow框架构建的Python库，专门用于图像分割任务。该库提供了高度封装的API接口，让开发者能够用极简的代码快速构建和部署图像分割神经网络模型。

核心特性

简洁高效的API设计：仅需2行代码即可创建神经网络模型
丰富的模型架构：包含4种主流分割网络架构
强大的预训练支持：提供25种不同的骨干网络(backbone)选择
预训练权重支持：所有骨干网络都提供ImageNet预训练权重

快速入门

基础模型创建

from segmentation_models import Unet

# 创建基础Unet模型
model = Unet()

进阶配置

# 使用ResNet34作为骨干网络，并加载ImageNet预训练权重
model = Unet('resnet34', encoder_weights='imagenet')

# 配置3分类任务，使用softmax激活函数
model = Unet('resnet34', classes=3, activation='softmax')

# 自定义输入形状(6通道输入)，不使用预训练权重
model = Unet('resnet34', input_shape=(None, None, 6), encoder_weights=None)

完整训练流程示例

from segmentation_models import Unet
from segmentation_models import get_preprocessing
from segmentation_models.losses import bce_jaccard_loss
from segmentation_models.metrics import iou_score

# 配置骨干网络
BACKBONE = 'resnet34'
preprocess_input = get_preprocessing(BACKBONE)

# 加载数据
x_train, y_train, x_val, y_val = load_data(...)

# 数据预处理
x_train = preprocess_input(x_train)
x_val = preprocess_input(x_val)

# 创建并编译模型
model = Unet(BACKBONE, encoder_weights='imagenet')
model.compile('Adam', loss=bce_jaccard_loss, metrics=[iou_score])

# 训练模型
model.fit(
    x=x_train,
    y=y_train,
    batch_size=16,
    epochs=100,
    validation_data=(x_val, y_val),
)

支持的模型架构

Unet：经典的编码器-解码器结构，特别适合医学图像分割
FPN (Feature Pyramid Network)：特征金字塔网络，擅长处理多尺度目标
Linknet：轻量级架构，计算效率高
PSPNet (Pyramid Scene Parsing Network)：金字塔场景解析网络，适合复杂场景理解

支持的骨干网络

该库支持多种主流CNN架构作为骨干网络：

| 类型 | 具体实现 | |--------------|--------------------------------------------------------------------------| | VGG | vgg16, vgg19 | | ResNet | resnet18到resnet152系列 | | SE-ResNet | seresnet18到seresnet152系列 | | ResNeXt | resnext50, resnext101 | | DenseNet | densenet121, densenet169, densenet201 | | Inception | inceptionv3, inceptionresnetv2 | | MobileNet | mobilenet, mobilenetv2 | | EfficientNet | efficientnetb0到efficientnetb5 |

所有骨干网络都提供在ImageNet数据集上预训练的权重。

高级技巧

模型微调策略

from segmentation_models import Unet
from segmentation_models.utils import set_trainable

# 初始阶段冻结编码器权重
model = Unet(backbone_name='resnet34', encoder_weights='imagenet', encoder_freeze=True)
model.compile('Adam', 'binary_crossentropy', ['binary_accuracy'])

# 仅训练解码器部分
model.fit(x, y, epochs=2)

# 解冻所有层继续训练
set_trainable(model)
model.fit(x, y, epochs=100)

处理非RGB数据

从头训练策略：

# 定义输入通道数N
N = x.shape[-1]
model = Unet(backbone_name='resnet34', encoder_weights=None, input_shape=(None, None, N))

通道转换策略：

from keras.layers import Input, Conv2D
from keras.models import Model

# 创建1x1卷积层将N通道映射到3通道
base_model = Unet(backbone_name='resnet34', encoder_weights='imagenet')
inp = Input(shape=(None, None, N))
l1 = Conv2D(3, (1, 1))(inp)
out = base_model(l1)
model = Model(inp, out, name=base_model.name)