5 行代码实现图像分割

163 阅读2分钟

原文链接 python1234.cn/archives/fi…

大家好,我是小寒。

图像分割是计算机视觉中的一个重要分支,可以用来解决许多计算机视觉问题,例如医学图像分析、背景编辑、自动驾驶和卫星图像分析等。

PixelLib 是一个 python 库,旨在使用几行 Python 代码轻松实现在图像和视频中进行实例分割。

新版本 PixelLib 提供了 PyTorch 作为后端,使用 PointRend 模型来执行实例分割

PointRend 是用于实现对象分割的最先进的神经网络。它生成准确的分割掩码并以高推理速度运行。

下载安装

PixelLib 支持 python 3.7 及以上版本

安装 PixelLib 及其依赖项
  • 安装 PyTorch

    PixelLib 支持这些 PyTorch 版本**(1.6.0、1.7.1、1.8.0和1.90)**。不支持 PyTorch 1.7.0,请勿使用任何低于 1.6.0 的 PyTorch 版本。

    conda install pytorch==1.9.0 torchvision==0.10.0 torchaudio==0.9.0 cudatoolkit=11.3 -c pytorch -c conda-forge
    
  • 安装 Pycocotools

    pip install pycocotools
    
  • 安装 PixelLib

    pip install pixellib
    

图像分割

PixelLib 使用五行 python 代码,实现在图像和视频中执行对象分割.

ins = instanceSegmentation()
ins.load_model("./data/image_seg/pointrend_resnet50.pkl")
ins.segmentImage("./data/image_seg/image1.jpg", show_bboxes=True, output_image_name="output_image1.jpg")

下面我们来看一下分割的结果。

原始图像:

分割后的图像:

从结果来看,效果还是不错的。

PixelLib 使执行实时对象分割成为可能,并增加了调整推理速度以适应实时预测的能力。

使用 4GB 容量的 Nvidia GPU 处理单个图像的默认推理速度约为0.20 秒

PixelLib 支持调速,调速模式有两种:

  • fast
  • rapid

1. Fast Mode

ins.load_model("pointrend_resnet50.pkl", detection_speed = "fast")

load_model 函数中,我们添加了参数 detection_speed 并将值设置为 fastfast 模式下单张图像处理时间达到 0.20秒

import pixellib
from pixellib.torchbackend.instance import instanceSegmentation

ins = instanceSegmentation()
ins.load_model("./data/image_seg/pointrend_resnet50.pkl")
ins.segmentImage("./data/image_seg/image1.jpg", show_bboxes=True, output_image_name="output_image1.jpg")

2. Rapid Mode

ins.load_model("pointrend_resnet50.pkl", detection_speed = "rapid")

load_model 函数中,我们添加了参数 detection_speed 并将值设置为rapid。 rapid 模式下单张图像处理时间可达 0.15秒。

图像分割中的自定义对象检测

使用的 PointRend 模型是在 COCO 数据集上进行预训练的,支持 80 类对象。

PixelLib 支持自定义对象检测,我们可以从支持的 80 类对象中选择符合我们目标的对象。

import pixellib
from pixellib.torchbackend.instance import instanceSegmentation

ins = instanceSegmentation()
ins.load_model("./data/image_seg/pointrend_resnet50.pkl")
target_classes = ins.select_target_classes(person = True)
ins.segmentImage("./data/image_seg/image1.jpg", show_bboxes=True, 
segment_target_classes = target_classes,output_image_name="output_image1.jpg")

调用函数 **select_target_classes **选择要分割的目标对象。

函数 segmentImage 有一个新参数 segment_target_classes 来从目标类中选择并根据它们过滤检测。

原始图像:

分割后的图像:

可以看到只检测到了图像中的人。