直接使用
请打开基于EasyCV的STDC图像语义分割示例,并点击右上角 “ 在DSW中打开” 。
EasyCV语义分割-STDC
本文将以STDC模型为例,介绍如何基于EasyCV进行语义分割模型训练和预测。
STDC(Short-Term Dense Concatenate network)是一种轻量化的语义分割模型,通过轻量化的STDC模块和边缘学习分支,在Cityscape数据集上能够达到71.9%的mIOU,同时在NVIDIA GTX 1080Ti上实现250FPS的速度。
运行环境要求
modelscope:tf1.15torch1.11-gpu-py37-cu113-ubuntu20.04 镜像, GPU机型 P100 or V100, 内存 32G
安装依赖包
1、安装mmcv-full
# 根据cuda和torch版本 ! pip install -U openmim && mim install mmcv-full==1.6.1
2、安装EasyCV算法包
可以通过如下命令安装,也可以拉取最新EasyCV代码进行安装
! echo y | pip uninstall pai-easycv && pip install http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/pkgs/whl/2023_02/pai_easycv-0.9.0-py3-none-any.whl ! echo y | pip uninstall mmdet
3、简单验证
from easycv.apis import *
图像语义分割模型训练&预测
下面示例将介绍如何基于EasyCV STDC算法训练语义分割模型,并进行预测。
数据准备
可以使用我们提供的示例数据快速走通流程,也可以使用自定义的数据集进行模型训练
! wget http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/data/small_cityscape_raw/small_cityscapes_raw.tar.gz && tar -zxf small_cityscapes_raw.tar.gz
该数据是提取部分cityscape数据,并将标签文件转换成了easycv所需的分割格式(标签文件是单通道的png图片,图片上的像素值即代表了label id)
文件目录如下
small_cityscapes_raw ├── imgs │ ├── train │ │ ├── aachen_000000_000019.png │ │ ├── aachen_000001_000019.png ... │ │ └── aachen_000173_000019.png │ └── val │ ├── bremen_000000_000019.png │ ├── bremen_000001_000019.png ... │ └── bremen_000315_000019.png └── labels ├── train │ ├── aachen_000000_000019.png │ ├── aachen_000001_000019.png ... │ └── aachen_000173_000019.png └── val ├── bremen_000000_000019.png ├── bremen_000001_000019.png ...
训练模型
下载示例配置文件, 进行STDC模型训练
! rm -rf stdc1_raw_8xb6_e1290.py ! wget http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/configs/segmentation/stdc/stdc1_raw_8xb6_e1290.py
为了适配小数据,我们对配置文件stdc1_raw_8xb6_e1290.py做如下字段的修改,减少训练epoch数目
如果使用自定义的数据集,则还需要修改config中数据集路径相关配置、以及类别数目和类别名配置
total_epochs = 10
!python -m torch.distributed.launch --nproc_per_node=1 \ /opt/conda/lib/python3.7/site-packages/easycv/tools/train.py \ stdc1_raw_8xb6_e1290.py --work_dir work_dir/example_seg --launcher pytorch --fp16
预测
查看训练产生的模型文件,其中_export.pt结尾的模型文件是最终导出用来做推理的模型
! ls work_dir/example_seg/*pt*
导入模型权重,并预测测试图片的检测结果,这里由于在小数据集上无法充分训练得到比较好的结果,这里提供在cityscape全量数据上训练的模型进行预测
! wget http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/segmentation/stdc/stdc1_cityscapes/epoch_1250_export.pth from easycv.predictors.segmentation import SegmentationPredictor output_ckpt = 'epoch_1250_export.pth' detector = SegmentationPredictor(output_ckpt) output = detector(['small_cityscapes_raw/imgs/val/bremen_000000_000019.png']) print(output) detector.show_result('small_cityscapes_raw/imgs/val/bremen_000000_000019.png', output[0]['seg_pred'], out_file='out.jpg')
from PIL import Image display(Image.open('out.jpg'))
将导出的模型文件上传至OSS,使用PAI-EAS中的easycv predictor即可部署在线服务