segmentation_models.pytorch

Segmentation Models PyTorch (SMP) 是一个基于 PyTorch 构建的 Python 图像语义分割库。它提供高层 API，分割模型就是标准的 torch.nn.Module，几行代码即可创建。 README 中描述的主要能力： - 12 种编码器-解码器架构：Unet、Unet++、MAnet、Linknet、FPN、PSPNet、PAN、DeepLabV3、DeepLabV3+、UPerNet、Segformer 和 DPT，每种都附有论文和文档链接。 - 800 多个预训练的卷积和 Transformer 编码器，包括对 timm 骨干网络的支持。编码器暴露中间特征，而不仅仅是最终分类特征，并附带预训练权重，旨在加速并稳定收敛。README 指出，轻量级选项（mobilenet/mobileone）适合低延迟或边缘推理，而更高容量的选项（convnext/swin/mit）适合复杂任务。 - 预处理辅助工具：get_preprocessing_fn 返回与编码器预训练（例如 ImageNet）匹配的预处理函数，README 称其可以改善结果和收敛，但在训练整个模型时并非强制要求。 - 分割训练中常用的损失函数和指标（Dice、Jaccard、Tversky 等）。 - 模型 API 选项：in_channels 用于任意输入通道数（文档说明了首个卷积权重的复用方式），aux_params 用于可选的辅助分类头（全局池化、可选 dropout、线性层、可选激活），在输出掩码的同时产生标签输出，encoder_depth 用于减少下采样阶段，使模型更轻量。 - 导出和部署友好：支持 ONNX 导出以及 torch script/trace/compile 兼容性。 示例包括二值分割（OxfordPets、CamVid）、多类分割（CamVid）的 notebook，使用预训练 Segformer、DPT 和 UPerNet 加载并运行推理，将模型保存/加载到本地或 Hugging Face Hub，以及导出为 ONNX。UPerNet、Segformer 和 DPT 的检查点在 Hugging Face 上有链接。 安装可通过 pip（segmentation-models-pytorch）或直接从 GitHub 进行。README 还提到了使用该库赢得竞赛的名人堂页面、一个使用 SMP 构建的项目（withoutBG 背景移除）、使用 make 目标进行开发安装、测试和代码检查的贡献说明、引用条目，以及许可：主要为 MIT，部分文件采用其他许可证，因此建议商业用户查看 LICENSES 文件和每个文件的声明。