用OpenMMLab来打CV比赛!

  • 2020 年 12 月 15 日
  • AI

Hi,大家吼呀!这里是糖葫芦喵喵~!

今天给大家带来的是Spark“数字人体”AI挑战赛——脊柱疾病智能诊断大赛 的方案思路。比赛传送门:

Spark“数字人体”AI挑战赛–脊柱疾病智能诊断大赛-天池大赛-阿里云天池tianchi.aliyun.com图标

在这个比赛中,我们将全程使用令人激动的OpenMMLab文末有获取开源代码地址的方法。

Part 0 OpenMMLab 简介

OpenMMLab 是香港中文大学与商汤科技为计算机视觉的不同方向建立统一而开放的代码库,包括各种计算机视觉问题中的主流算法。该代码库采用模块化设计,使用十分方便。

mmdetection举例,我们需要关注的代码主要集中于:

  • configs/中对各种模型、训练策略、数据集的配置。
  • mmdet/datasets/内置主要数据集(coco\voc\cityscapes等),并支持自定义数据集。
  • mmdet/datasets/pipelines内置主要数据处理过程的操作。
  • mmdet/models/内置主要模型与损失函数的实现方法。
  • mmdet/apis/内置训练测试的API接口。
  • tools/内置训练测试的脚本。

通过修改上述内容我们可以快速的将自己的数据集应用于这个强大的框架。下面我们来通过一个比赛实践一下吧!

Part 1 任务描述

本次挑战的目标是通过人工智能算法自动分析核磁共振影像来检测和分类脊柱的退行性改变,包括正常椎体、退变椎体、正常椎间盘、椎间盘突出、膨出等特征。参与者需要提供全自动算法来定位椎体和椎间盘的位置和相应分类。

椎体:6个位置(T12,L1,L2,L3,L4,L5),2种病灶(V1,V2)

椎间盘:6个位置(T12-L1,L1-L2,L2-L3,L3-L4,L4-L5,L5-S1),5种病灶(V1,V2,V3,V4,V5)。

评价指标:F1

Part 2 数据分析

赛题提供了150张训练集以及51张测试集,格式为DICOM。DICOM影像包含T1和T2 矢状面影像,以及T2 轴状位影像(FSE/TSE):

已标注的数据位于T2矢状位的中间帧。从图中我们可以看出,定位信息包含在T2矢状位中,而椎体\椎间盘对应的轴状位可以用来辅助进行病灶分类。

Part 3 脊椎框检测-mmdetection

我们选择top-down方式进行关键点检测。类似于人体关键点检测首先要找到人,这里我们需要检测出脊椎框。

首先我们要构造coco格式标注:

images:

  • file_name:图像文件路径
  • height:图像高
  • width:图像宽
  • id:唯一标识图像ID

annotations:

  • bbox:xywh检测框
  • category_id:该检测框类别ID
  • image_id:该检测框所属图像ID (对应images中的id)
  • id:全局唯一标识该检测框ID
  • area:bbox面积
  • keypoints:xyv关键点序列(用于关键点检测)
  • num_keypoints:关键点数量(用于关键点检测)
  • iscrowd:是否是单独物体(用于语义分割
  • segmentation:物体掩码(用于语义分割)

categories:

  • id:类别ID
  • name:类别ID对应的类别名称

我们通过计算所有标注坐标的最大最小值为每一张T2矢状位构建脊椎框。因为我们要检测的只有脊椎框,所以我们的categories只有一个:”categories”: [{“id”: 0, “name”: “box”}]

标注生成好之后我们只需要简单更改配置文件:

修改config/_base_/model/对应模型bbox_head的num_classes为1

在config/_base_/dataset/coco_detection.py中为data指定classes = (‘box’)

修改mmdet/core/evaluation/class_names.py中coco_classes return [‘box’,]。

修改mmdet/datasets/coco.py中CLASSES = (‘box’, )。

然后就可以愉快的训练了!

# 训练
!./mmdetection/tools/dist_train.sh ./mmdetection/configs/faster_rcnn/faster_rcnn_x101_64x4d_fpn_1x_coco.py 2
# 测试
!python ./mmdetection/tools/test.py ./mmdetection/configs/faster_rcnn/faster_rcnn_x101_64x4d_fpn_1x_coco.py ./mmdetection/work_dirs/x/epoch_80.pth --out ../data/results_bbox_test.pkl 

测试结果:

Part 4 关键点检测-mmpose

有了脊椎框我们就可以开始愉快的进行关键点检测了。依旧构建coco标注格式:

images:

  • file_name:图像文件路径
  • height:图像高
  • width:图像宽
  • id:唯一标识图像ID

annotations:

  • bbox:xywh检测框
  • category_id:该检测框类别ID
  • image_id:该检测框所属图像ID (对应images中的id)
  • id:全局唯一标识该检测框ID
  • area:bbox面积
  • keypoints:[x1,y1,v1,x2,y2,v2…]其中v表示是否可见2/遮挡1/未标注0
  • num_keypoints:关键点数量(用于关键点检测)
  • iscrowd:是否是单独物体(用于语义分割)
  • segmentation:物体掩码(用于语义分割)

categories:

  • id:(检测框)类别ID
  • name:(检测框)类别名称
  • supercategory:(检测框)类别名称
  • keypoints:关键点类别名称
  • skeleton:关键点之间的连通性

这里我们定义了 12种位置以及对应的连通关系:

“keypoints”: [“T12″,”L1″,”L2″,”L3″,”L4″,”L5″,”T12-L1″,”L1-L2″,”L2-L3″,”L3-L4″,”L4-L5″,”L5-S1”]

“skeleton”: [[0,6],[6,1],[1,7],[7,2],[2,8],[8,3],[3,9],[9,4],[4,10],[10,5],[5,11]]

标注生成好之后我们只需要简单更改配置文件:

修改configs/top_down对应channel_cfg的关键点数量;取消TopDownRandomFlip/TopDownHalfBodyTransform;flip_test设置为False。

类似的,我们可以为验证集pipeline添加TopDownGenerateTarget,然后修改models/detectors/top_down.py中的forward_test方法来计算验证集的acc:

        if target is not None:
            _, avg_acc, cnt = pose_pck_accuracy(
                output[target_weight.squeeze(-1) > 0].unsqueeze(
                    0).detach().cpu().numpy(),
                target[target_weight.squeeze(-1) > 0].unsqueeze(
                    0).detach().cpu().numpy())

然后就可以愉快的训练了!

# 训练
!./mmpose/tools/dist_train.sh ./mmpose/configs/top_down/hrnet/coco/hrnet_w32_coco_384x288.py 2 

测试过程中,我们首先用Part3中训练的模型计算测试集的脊椎框,然后将图片和对应的脊椎框送入关键点检测模型中即可。

测试结果:

椎体和椎间盘分开看:

Part 5 病灶分类-mmclassification

到此为止我们将椎体无脑设为v2椎间盘设为v1提交就有0.4+的成绩,说明我们的定位还是比较准的。赛题的难点是在病灶分类上。一个简单的手段就是对关键点附近的图像进行裁剪,然后使用分类算法进行分类。

这里我们构建imagenet格式数据:

  • train目录下按照类别ID设置子目录,存放对应图像
  • val目录下存放全部验证集
  • meta/val.txt存放验证集与类别ID对应信息

修改config/_base_/model/对应模型的num_classes为7

然后就可以愉快的训练了!

!python -u tools/train.py configs/imagenet/seresnext50_32x4d_batch256.py

验证集上可以达到0.6+的准确率。相信大家调参后能得到更好的结果。我们的全部代码将在后续整理完毕后开源。

是不是很简单?OpenMMLab,永远滴神!

关注微信公众号AILIGHT并回复:lumbar 获取开源代码地址链接

还没有关注AILIGHT的朋友,扫描下方二维码识别关注吧。有什么建议也可以后台留言哦!