pytorch semantic segmentation using torchvision

最新推荐文章于 2024-08-18 10:45:14 发布

原创最新推荐文章于 2024-08-18 10:45:14 发布 · 335 阅读

0 ·

CC 4.0 BY-SA版权

pytorch 专栏收录该内容

28 篇文章

订阅专栏

本文分享了如何使用预训练的FCN和DeepLabV3模型进行图像分割，通过代码实例展示了将标签映射到颜色的过程，并介绍了在`torchvision`中加载和使用模型的步骤。适合理解深度学习在图像识别中的应用。

一些有用的小的代码片段。

# use the following function to convert this 2D image to an RGB image where each label is mapped to its 
# corresponding color.
# 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()

from torchvision import models
fcn = models.segmentation.fcn_resnet101(pretrained=True).eval()
!wget -nv https://www.learnopencv.com/wp-content/uploads/2021/01/horse-segmentation.jpeg -O horse.png
segment(fcn, './horse.png')
dlab = models.segmentation.deeplabv3_resnet101(pretrained=1).eval()
segment(dlab, './horse.png')