7 MMSegmentation 代码教学

news2024/9/28 11:24:58

本文是openmmlab AI实战营的第七次课程的笔记，以下是我比较关注的部分。

本次课程的主要内容是 mmsegmentation的代码教学实战，我会稍微详细记载。

环境安装

云在线平台：Featurize

推荐代码运行环境：GPU RTX3060 ,CUDA 11.2

安装pytorch

!pip3 install install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio==0.10.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html

用mim 安装mmcv

pip install -U openmim
mim install  mmengine
mim install 'mmcv==2.0.0rc3'

安装其它工具包

!pip install opencv-python pillow matplotlib seaborn tqdm 'mmdet>=3.0.0rc1' -i https://pypi.tuna.tsinghua.edu.cn/simple

下载 MMSegmentation

# 删掉原有的 mmtracking 文件夹（如有）
!rm -rf mmsegmentation

# 从 github 上下载最新的 mmsegmentation 源代码
!git clone https://github.com/open-mmlab/mmsegmentation.git -b dev-1.x

安装 MMSegmentation

!pip install -v -e .

下载预训练模型权重文件和视频素材

import os
# 创建 checkpoint 文件夹，用于存放预训练模型权重文件
os.mkdir('checkpoint')
# 创建 outputs 文件夹，用于存放预测结果
os.mkdir('outputs')
# 创建 data 文件夹，用于存放图片和视频素材
os.mkdir('data')

下载预训练模型权重

Model Zoo：https://github.com/open-mmlab/mmsegmentation/blob/master/docs/en/model_zoo.md

# 从 Model Zoo 获取 PSPNet 预训练模型，下载并保存在 checkpoint 文件夹中
!wget https://download.openmmlab.com/mmsegmentation/v0.5/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-296659

下载素材

如果报错Unable to establish SSL connection.，重新运行代码块即可。

# 伦敦街景图片
!wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220713-mmdetection/images/street_uk.jpeg -P data

# 上海驾车街景视频，视频来源：https://www.youtube.com/watch?v=ll8TgCZ0plk
!wget https://zihao-download.obs.cn-east-3.myhuaweicloud.com/detectron2/traffic.mp4 -P data

# 街拍视频，2022年3月30日
!wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220713-mmdetection/images/street_20220330_174028.mp4 -P data

检查安装成功

# 检查 Pytorch
import torch,torchvision
print('Pytorch 版本',torch.__version__)
print('CUDA 是否可用',torch.cuda.is_available())

Pytorch 版本 1.10.1+cu113

CUDA 是否可用 True

# 检查 mmcv
import mmcv from mmcv.opsimportget_compiling_cuda_version,get_compiler_version
print('MMCV版本',mmcv.__version__)
print('CUDA版本',get_compiling_cuda_version())
print('编译器版本',get_compiler_version())

MMCV版本 2.0.0rc3

CUDA版本 11.3

编译器版本 GCC 9.3

# 检查 mmsegmentation
import mmseg from mmseg.utilsimportregister_all_modules
from mmseg.apis import inference_model,init_model
print('mmsegmentation版本',mmseg.__version__)

mmsegmentation版本 1.0.0rc4

图像分割

import os 
os.chdir('mmsegmentation')

载入测试图像

from PIL import Image

# Image.open('data/street_uk.jpeg')

图像预测-命令行

PSPNet语义分割算法

!python demo/image_demo.py \data/street_uk.jpeg \
        configs/pspnet/pspnet_r50-d8_4xb2-40k_cityscapes-512x1024.py \
        https://download.openmmlab.com/mmsegmentation/v0.5/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth \
        --out-fileoutputs/B1_uk_pspnet.jpg \
        --devicecuda:0 \
        --opacity0.5

Image.open('outputs/B1_uk_pspnet.jpg')

视频分割

进入 mmsegmentation 主目录

import os
os.chdir('mmsegmentation')

视频预测-命令行

!python demo/video_demo.py \data/traffic.mp4 \
        configs/pspnet/pspnet_r50-d8_4xb2-40k_cityscapes-512x1024.py \
        https://download.openmmlab.com/mmsegmentation/v0.5/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth \
        --devicecuda:0 \
        --output-fileoutputs/B3_video.mp4 \
        --opacity0.5

在自己的数据集上训练语义分割模型

进入 mmsegmentation 主目录

import os
os.chdir('mmsegmentation')

导入工具包

import numpy as np
from PIL import Image

import os.path as osp
from tqdm import tqdm

import mmcv
import mmengine
import matplotlib.pyplot as plt
%matplotlib inline

数据集下载和预处理

Standord Background 数据集主页：http://dags.stanford.edu/projects/scenedataset.html

该数据集包含 715 张由 LabelMe, MSRC, PASCAL VOC 和 Geometric Context 四个开源数据集中挑选出来的图片。

主要为室外场景，图像尺寸为320x240。

共 8 个类别：

天空 sky

树 tree

道路 road

草 grass

水 water

建筑物 building

山 mountain

前景物体

原版数据集下载及预处理（略）

# 数据集图片和标注路径
data_root = 'iccv09Data'
img_dir = 'images'
ann_dir = 'labels'

# 类别和对应的颜色
classes = ('sky', 'tree', 'road', 'grass', 'water', 'bldg', 'mntn', 'fg obj')
palette = [[128, 128, 128], [129, 127, 38], [120, 69, 125], [53, 125, 34], 
           [0, 11, 123], [118, 20, 12], [122, 81, 25], [241, 134, 51]]

直接下载子豪兄预处理好的数据集

!wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20230130-mmseg/dataset/iccv09Data.zip

# 解压
!unzip iccv09Data.zip >> /dev/null
# 删除压缩包
!rm -rf iccv09Data.zip

查看数据集中的图片

Image.open('iccv09Data/images/6000124.jpg')

查看语义分割标注 Segmentation Map

Image.open('iccv09Data/labels/6000124.png')

In [12]:

import matplotlib.patches as mpatches
img = Image.open('iccv09Data/labels/6000124.png')
plt.figure(figsize=(8, 6))
im = plt.imshow(np.array(img.convert('RGB')))

# 图例小块
patches = [mpatches.Patch(color=np.array(palette[i])/255., label=classes[i]) for i in range(8)]
# 图例
plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0., fontsize='large')
plt.show()

修改数据集类

After downloading the data, we need to implement load_annotations function in the new dataset class StanfordBackgroundDataset.

from mmseg.registry import DATASETS
from mmseg.datasets import BaseSegDataset

@DATASETS.register_module()
class StanfordBackgroundDataset(BaseSegDataset):
  METAINFO = dict(classes = classes, palette = palette)
  def __init__(self, **kwargs):
    super().__init__(img_suffix='.jpg', seg_map_suffix='.png', **kwargs)

文档：https://github.com/open-mmlab/mmsegmentation/blob/master/docs/en/tutorials/customize_datasets.md#customize-datasets-by-reorganizing-data

修改config配置文件

# 下载 config 文件 和 预训练模型checkpoint权重文件
!mim download mmsegmentation --config pspnet_r50-d8_4xb2-40k_cityscapes-512x1024 --dest .

processing pspnet_r50-d8_4xb2-40k_cityscapes-512x1024...

# modify num classes of the model in decode/auxiliary head
cfg.model.decode_head.num_classes = 8
cfg.model.auxiliary_head.num_classes = 8
# 修改数据集的 type 和 root

查看完整config配置文件

print(cfg.pretty_text)

准备训练

from mmengine.runner import Runner
from mmseg.utils import register_all_modules

# register all modules in mmseg into the registries
# do not init the default scope here because it will be init in the runner
register_all_modules(init_default_scope=False)
runner = Runner.from_cfg(cfg)

开始训练

如果遇到报错CUDA out of memeory，重启实例或使用显存更高的实例即可。

runner.train()

使用训练得到的模型预测图像

from mmseg.apis import init_model,inference_model,show_result_pyplot

# 初始化模型
checkpoint_path='./work_dirs/tutorial/iter_800.pth'
model=init_model(cfg,checkpoint_path,'cuda:0')

Loads checkpoint by local backend from path: ./work_dirs/tutorial/iter_800.pth

# 载入测试图像
img=mmcv.imread('iccv09Data/images/6000124.jpg')
# img = mmcv.imread('iccv09Data/images/6000295.jpg')

# 预测
result=inference_model(model,img)

# 可视化预测结果
visualization = show_result_pyplot(model, img, result, opacity=0.8, title='MMSeg', out_file='outputs/C.jpg')
plt.imshow(mmcv.bgr2rgb(visualization))
plt.show()