### 关于 `segmentation_models.pytorch` 库中的 UNet++ 模型 #### 使用说明 `segmentation_models.pytorch` 提供了多种语义分割网络架构，其中包括改进版本的 U-Net 架构——UNet++。此模型通过引入密集跳跃连接来增强特征重用和梯度传播效率[^1]。 为了初始化一个带有预训练权重（例如 ImageNet 上面训练得到）的 UNet++ 实例，可以按照如下方式进行： ```python import segmentation_models_pytorch as smp model = smp.UnetPlusPlus( encoder_name="resnet34", # choose encoder, e.g. resnet34, efficientnet-b0 encoder_weights="imagenet", # use `imagenet` pre-trained weights for encoder initialization in_channels=3, # model input channels (default 3) classes=1, # model output channels (number of classes in your dataset) ) ``` 上述代码片段展示了如何利用 ResNet34 编码器以及其在 ImageNet 数据集上的预训练参数快速搭建起一个用于二元分类任务的 UNet++ 网络结构[^2]。 对于多类别分割问题，则需调整最后一层输出通道数以匹配目标数据集中类别的数量: ```python model = smp.UnetPlusPlus( encoder_name="efficientnet-b0", encoder_weights="imagenet", in_channels=3, classes=8, # 假设有8个不同的物体类别加背景一共9种情况 ) ``` 这里选择了 EfficientNet-B0 作为编码部分的基础，并设置了输入图片的颜色通道数目为三（RGB），同时指定最终预测结果应包含八个前景对象加上一个默认背景标签共九维向量表示法[^4]。 #### 示例应用 下面给出一段简单的例子，演示怎样加载已有的 UNet++ 模型并执行推理操作： ```python from PIL import Image import torch from torchvision.transforms.functional import to_tensor # 加载之前保存好的 checkpoint 或者直接定义新模型 checkpoint_path = 'path_to_your_checkpoint.pth' device = "cuda" if torch.cuda.is_available() else "cpu" model.load_state_dict(torch.load(checkpoint_path)) model.to(device) def predict(image_file): img = Image.open(image_file).convert('RGB') tensor_img = to_tensor(img).unsqueeze(0) # 添加批次维度 tensor_img = tensor_img.to(device) with torch.no_grad(): pred_mask = model(tensor_img)[0].cpu().numpy() return pred_mask.argmax(axis=0) # 获取每像素最可能属于哪个类别 if __name__ == '__main__': prediction = predict('./example_image.png') print(prediction.shape) # 输出形状应该与原始图像相同 ``` 这段脚本首先读取一张测试用图象文件，将其转换成张量形式送入 GPU 进行前向计算获得掩膜矩阵；最后返回每个位置上概率最高的那一种类别编号组成的二维数组。