使用python基于DeepLabv3实现对图片进行语义分割

news2024/10/4 14:54:25

DeepLabv3 介绍

DeepLabv3 是一种先进的语义分割模型，由 Google Research 团队提出。它在 DeepLab 系列模型的基础上进行了改进，旨在提高图像中像素级分类的准确性。以下是 DeepLabv3 的详细介绍：

概述DeepLabv3 是 DeepLab 系列中的第三代模型，专门用于解决语义分割任务。语义分割是指将图像中的每个像素分配给特定类别的过程，例如将图像中的不同对象（如人、车、建筑物等）进行精确划分。
主要特点
- 空洞卷积（Atrous Convolution）：DeepLabv3 使用空洞卷积来增加感受野，同时保持特征图的空间分辨率。空洞卷积通过在标准卷积核中插入“空洞”来扩大卷积核的有效大小，从而捕捉更大范围的上下文信息。
- 多尺度输出（Multi-Scale Output）：为了处理不同尺度的对象，DeepLabv3 在多个不同的空洞率下应用空洞卷积，并将这些不同尺度的特征图融合在一起。这种多尺度特征融合有助于提高对小物体和大物体的识别能力。
- 空间金字塔池化（Atrous Spatial Pyramid Pooling, ASPP）：ASPP 模块是 DeepLabv3 的核心组件之一，它通过使用不同空洞率的空洞卷积以及全局平均池化来捕获多尺度信息。这有助于模型更好地理解图像中的不同尺度和结构。
- 解码器模块（Decoder Module）：DeepLabv3 引入了一个简单的解码器模块，用于恢复高分辨率的分割结果。解码器模块通过上采样操作将低分辨率的特征图恢复到原始图像的分辨率，并结合低层次的特征图以增强细节信息。
架构DeepLabv3 的架构可以分为以下几个主要部分：
- 主干网络（Backbone Network）：通常使用预训练的深度卷积神经网络（如 ResNet 或 Xception）作为主干网络，提取高层次的特征表示。
- ASPP 模块：在主干网络的输出特征图上应用 ASPP 模块，生成多尺度的特征表示。•解码器模块：通过上采样操作将 ASPP 输出的特征图恢复到原始图像的分辨率，并结合低层次的特征图以增强细节信息。
- 最终分类层：在解码器模块的输出上应用一个逐像素的分类层，生成最终的分割结果。
训练与推理•损失函数：通常使用交叉熵损失函数或其变体（如加权交叉熵、Dice 损失等）来训练模型。
- 数据增强：为了提高模型的泛化能力，可以采用各种数据增强技术，如随机裁剪、翻转、缩放等。
- 推理：在推理阶段，输入图像经过模型处理后，生成每个像素的类别标签，形成最终的分割结果。

代码实现

import torch
import torch.nn.functional as F
from torchvision.models.segmentation import deeplabv3_resnet50
from torchvision import models
from PIL import Image
from torchvision import transforms as T
import numpy as np


# 反归一化
def denormalize(tensor, mean, std):
    # 创建一个新的变换来反归一化
    denormalize_transform = T.Normalize(mean=[-m/s for m, s in zip(mean, std)], std=[1/s for s in std])
    return denormalize_transform(tensor)

if __name__ == '__main__':
    # 加载预训练的DeepLabv3模型
    model = deeplabv3_resnet50(weights=models.segmentation.DeepLabV3_ResNet50_Weights.COCO_WITH_VOC_LABELS_V1,
                               num_classes=21)
    model.eval()
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    # 将图片转换为模型需要的输入格式
    transform = T.Compose([
        T.Resize(1024),
        T.ToTensor(),
        T.Normalize(mean=mean, std=std),
    ])

    # 加载图片
    # image = Image.open("fenge.jpg")
    image = Image.open("测试2.jpg")
    image_tensor = transform(image).unsqueeze(0)

    # 进行推理
    with torch.no_grad():
        output = model(image_tensor)['out']
        output_predictions = F.softmax(output, dim=1).argmax(dim=1)

    # 保存结果
    print(output_predictions.shape)
    print('----')
    output_predictions = output_predictions.squeeze(0).cpu().numpy()
   
    # 使用Pillow创建图像
    print(output_predictions)
    # 将numpy数组转换为PIL图像
    img = Image.fromarray(np.uint8(output_predictions * 255))  # 将数据缩放到0-255范围内

    # 保存图像
    img.save('output.png')