Pytorch实现图像语义分割(初体验)

本文介绍了使用Pytorch进行图像语义分割的初步实践,通过FCN模型实例展示了如何加载预训练模型、进行图像预处理和预测,并分享了部分测试结果。读者可通过实际操作体验这一过程。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

Pytorch实现图像语义分割(初体验)

这些天在学习图像语义分割相关的知识,并简单写了篇概述。原本想先看几篇经典论文,如全卷积网络FCN,奈何英语水平有限,翻译起来实在费劲。想来不如先直接体验一下语义分割的效果,果然实践起来还挺有趣的。遂将过程记录如下。

代码实现

from torchvision import models
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torchvision.transforms as T
import numpy as np


# Define the helper function
def decode_segmap(image, nc=21):
    label_colors = np.array([(0, 0, 0),  # 0=background
                             # 1=aeroplane, 2=bicycle, 3=bird, 4=boat, 5=bottle
                             (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128),
                             # 6=bus, 7=car, 8=cat, 9=chair, 10=cow
                             (0, 128, 128), (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0),
                             # 11=dining table, 12=dog, 13=horse, 14=motorbike, 15=person
                             (192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (192, 128, 128),
                             # 16=potted plant, 17=sheep, 18=sofa, 19=train, 20=tv/monitor
                             (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128)])

    r = np.zeros_like(image).astype(np.uint8)
    g = np.zeros_like(image).astype(np.uint8)
    b = np.zeros_like(image).astype(np.uint8)

    for l in range(0, nc):
        idx = image == l
        r[idx] = label_colors[l, 0]
        g[idx] = label_colors[l, 1]
        b[idx] = label_colors[l, 2]

    rgb = np.stack([r, g, b], axis=2)
    return rgb


def segment(net, path):
    img = Image.open(path)
    plt.imshow(img)
    plt.axis('off')
    plt.show()
    # Comment the Resize and CenterCrop for better inference results
    trf = T.Compose([T.Resize(256),
                     T.CenterCrop(224),
                     T.ToTensor(),
                     T.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])])
    inp = trf(img).unsqueeze(0)
    out = net(inp)['out']
    om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()
    rgb = decode_segmap(om)
    plt.imshow(rgb)
    plt.axis('off')
    plt.show()


fcn = models.segmentation.fcn_resnet101(pretrained=True).eval()
# dlb = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()

girl = '../img/girl_dog.jpg'
segment(fcn, girl)
# segment(dlb, girl)

参考链接:https://learnopencv.com/pytorch-for-beginners-semantic-segmentation-using-torchvision/
代码整体理解相对比较简单,详细内容在参考链接中讲解得很清除,我也不必再做赘述。

测试结果

下面展示部分代码运行结果。







可能图像分割的效果不是那么得好,但整体而言还是实现了语义分割,大家也可以自己找一些图片进行测试(注意找的图片要求是label_colors中的),如对代码有疑问可留言交流。

评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值