CenterNet代码讲解

CenterNet电子书

  • 公众号:GiantPandaCV
  • 电子书《CenterNet源码解析》
  • 基于非官方的CenterNet实现代码 更适合阅读和理解,dataloader、hourglass、训练流程更加的简单
  • 使用VOC 数据集

1、环境配置

  • 环境要求

    1. python3.5
    2. pytorch==0.4 or 1.1.0 or 1.0.0
    3. tensorboardX(可选)
  • 配置

    1. 将cudnn的 batch norm关闭。
      1. 打开torch/nn/functional.py文件,找到torch.batch_norm这一行,将torch.backends.cudnn.enabled选项更改为False
    2. 克隆项目
      CenterNet_ROOT=/path/to/clone/CenterNet
      git clone https://github.com/zzzxxxttt/pytorch_simple_CenterNet_45 $CenterNet_ROOT
      
    3. 安装cocoAPI
      cd $CenterNet_ROOT/lib/cocoapi/PythonAPI
      make 
      python setup.py install --user
      
    4. 编译可变形卷积DCN
      1. 在pyTorch0.4.1中,将$CenterNet_ROOT/lib/DCNv2_old 复制为 $CenterNet_ROOT/lib/DCNv2
      2. 在pyTorch1.1.0 or pyTorch1.0.0中,将$CenterNet_ROOT/lib/DCNv2_new复制为 $CenterNet_ROOT/lib/DCNv2
      3. 然后开始编译
        cd $CenterNet_ROOT/lib/DCNv2
         ./make.sh
        
    5. 编译NMS
      cd $CenterNet_ROOT/lib/nms
      make
      
    6. 对于COCO数据集 。将annotations,train2017,val2017,test2017放在 $CenterNet_ROOT/data/coco
      • 下载地址:http://cocodataset.org/#download
    7. 对于Pascal VOC数据集。下载以后将 annotations, images, VOCdevkit 放在 $CenterNet_ROOT/data/voc
      • 网盘地址:https://pan.baidu.com/share/init?surl=z6BtsKPHh2MnbfT25Y4wYw 密码:4iu2
    8. 如果选择Hourglass-104作为骨干网络,下载CornerNet预训练模型。将下载的权重checkpoint.t7 放到 $CenterNet_ROOT/ckpt/pretrain 中
      • 百度网盘链接:https://pan.baidu.com/s/1tp9-5CAGwsX3VUSdV276Fg 密码:y1z4
  • 配置自己的数据集

    • 作者习惯Pascal VOC数据集,以Pascal VOC数据集为例,修改自己的数据集。假设笔者只有一个类,“person”,按照这一个类进行修改,其他的类别也容易修改
    1. VOC类别修改
      • 将datasets/pascal.py 中的第16行内容
      VOC_NAMES = ['__background__', "aeroplane", "bicycle", 
      "bird", "boat","bottle", "bus", "car", "cat", "chair", "cow", "diningtable","dog","horse", 
      "motorbike", "person", "pottedplant", "sheep", "sofa","train", "tvmonitor"]
      
      • 修改为自己的名称
      VOC_NAMES = ['__background__', 'person']
      
      • 将datasets/pascal.py中第33行内容
      num_classes=20 修改为自己对应的类别个数 num_classes=1
      
      • 将 datasets/pascal.py 中的第 35 行内容:
      self.valid_ids = np.arange(1, 21, dtype=np.int32) 中的 21 修改为 2(只有一个类,在原来基础上加一)
      
    2. annotations
      • VOC数据集中没有annotations中所需要的json文件,这部分需要重新构建
        下面是一个VOC转COCO格式的脚本,需要修改xml pathJson file的名称

      • 注意这里json文件的命名要通过datasets/pascal.py中第44到48行的内容确定的

    self.split = split
    self.data_dir = os.path.join(data_dir, 'voc')
    self.img_dir = os.path.join(self.data_dir, 'images')
    _ann_name = {'train': 'trainval0712', 'val': 'test2007'}
    self.annot_path = os.path.join(self.data_dir, 'annotations', 'pascal_%s.json' % _ann_name[split])
    
    -   这里笔者为了方便命名对这些字段进行了修改
    
    self.data_dir = os.path.join(data_dir, 'voc') # ./data/voc
    self.img_dir = os.path.join(self.data_dir, 'images') # ./data/voc/images
    _ann_name = {'train': 'train2020', 'val': 'test2020'}
    # 意思是需要 json 格式数据集
    self.annot_path = os.path.join(
    self.data_dir, 'annotations', 'pascal_%s.json' % _ann_name[split])
    
    -  所以要求json的命名可以按照以下格式准备
    
    # ./data/voc/annotations
        # - pascal_train2020
        # - pascal_test2020
    
    -   数据集总体格式为
    
    - data
        - voc
            - annotations
                - pascal_train2020.json
                - pascal_test2020.json
            - images
                - *.jpg
        - VOCdevkit(这 个 文 件 夹 主 要 是 用 于 测 评)
            - VOC2007
                - Annotations
                    - *.xml
                - JPEGImages
                    - *.jpg
                - ImageSets
                    - Main
                        - train.txt
                        - val.txt
                        - trainval.txt
                        - test.txt
    
    1. 其他
      • 在datasets/pascal.py中21-22行, 标准差和方差最好替换成自己数据集的标准差和方差
      VOC_MEAN = [0.485, 0.456, 0.406]
      VOC_STD = [0.229, 0.224, 0.225]
      
  • 训练和测试

    1. 训练命令 训练命令比较多,可以写一个shell脚本来完成
      python train.py --log_name pascal_resdcn18_384_dp \
                      --dataset pascal \
                      --arch resdcn_18 \
                      --img_size 384 \
                      --lr 1.25e-4 \
                      --lr_step 45,60 \
                      --batch_size 32 \
                      --num_epochs 70 \
                      --num_workers 10
      
      • log name 表示记录日志的名称
      • dataset使用的数据集,pascal表示使用pascal voc格式
      • arch代表选择的backbone类型,有以下几种
        • large_hourglass
          • small_hourglass
          • resdcn_18
          • resdcn_34
          • resdcn_50
          • resdcn_101
          • resdcn_152
      • img_size 控制图片的长和宽
      • lr和lr_step 控制学习率大小及变化
      • batch_size 一个批次处理图片的数量
      • num epochs 代表学习数据集的总次数
      • num workers 代表开启多少个线程来加载数据集
    2. 测试命令 需要注意img_size要和训练的一致
      python test.py --log_name pascal_resdcn18_384_dp \
                  --dataset pascal \
                  --arch resdcn_18 \
                  --img_size 384
      
      
      • flip test 属于 TTA(Test Time Augmentation),可以一定程度上提高 mAP。
      # flip test
      python test.py --log_name pascal_resdcn18_384_dp \
                      --dataset pascal \
                      --arch resdcn_18 \
                      --img_size 384 \
                      --test_flip
      
      
  • 结果

    • 每隔 5 个 epoch 将进行一次 eval,在自己的数据集上最终可以得到 90% 左右的 mAP。
    • 笔者将已经改好的单类的CenterNet放在Github上:https://github.com/pprp/SimpleCVReproduction/tree/master/CenterNet

二、数据集加载过程

  • CenterNet如何加载数据,将标注数据转换为高斯分布的形式
  1. Yolov3 和 CenterNet流程对比
    • Yolov3
      1. 对图片进行resize,长宽都是32的倍数
      2. 图片经过特征提取后,空间分辨率变为原来的1/32
      3. 得到代表不同尺度的特征框,目标框表示为(x,y,w,h,c),分别代表左上角坐标,宽和高,含有物体的置信度
      4. 训练完成之后,使用非极大值抑制算法得到最终的目标框
    • CenterNet
      1. 对图片进行resize,长和宽一般相等,并且至少为4的倍数
      2. 图片经过网络的特征提取后,得到的特征图分辨率依旧较大,是原来的1/4。
        • 这是因为CenterNet采用类似人体姿态估计网络,需要的分辨率比较大
      3. 训练过程中,CenterNet得到一个heatmap,所以加载标签的时候,需要转换为类似的heatmap热图
      4. 测试过程中,只需要从热图中提取目标,这样就不需要使用NMS,降低了计算量
  2. CenterNet部分详解
    • 输入的图片(W,H,3),长宽和三个通道。
    • 输出的热图Y大小(W/R, H/R, C) ,R表示缩小的倍数, C表示类别数
    • 在COCO数据集中,R=4,C=80。 因为COCO有80个类
    • 如果Y(x,y,c)为1,代表在(x,y)的位置,检测到c类的目标
    • ground truth 也必须是热图的形式。标注的内容包含(x1,y1,x2,y2,c)即目标左上角坐标,右下角坐标和类别c, 按照以下流程转换ground truth
      1. 得到原图中心坐标 p = ((x1+x2)/2,(y1+y2)/2) 获得中心坐标
      2. 得到下采样后的feature map中对应的中心目标 p~= p/R 向下取整
      3. 如果输入图像时512,那么feature map的分辨率为128,将标注的目标以高斯核的方式将关键点映射到特征图上
        [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-WeHMytGf-1605335260862)(ABEF722B63F7430EAB7548C495929623)]
  3. 代码部分(如何实现高斯分布)
    • datasets/pascal.py 代码主要从getitem函数入手
    • gaussian_radius函数,如何获得高斯半径
    • draw_umich_guassian函数,将高斯分布分散到heatmap上
  4. heatmap上应用高斯核
    • CenterNet实在CornerNet的基础上进行改动,有很多原始代码
    • 当绿框和红框的IOU>0.7时,那么认为绿框可以被接受,反之不能接受
    • 问题如何确定高斯半径r,让绿框和红框的IOU大于0.7
      • r的取值转换为一个一元二次方程有解问题,r的取值可以通过求根公式获得
    • 在原始的版本的,求高斯半径r有bug,有人修正了这个bug之后,可以让AR提升1-3个百分点
  5. 高斯分布添加到heatmap上
    • 具体代码在utils/image/def gaussian2D之中
    • 在utils/image/def draw_umich_gaussian 之中

3、骨干网络之hourglass

  • CenterNet中主要提供了三个骨干网络ResNet-18(ResNet-101),DLA-34,Hourglass-104
  1. Ground Truth Heatmap
    • CenterNet为什么要沿用CornerNet的半径计算方式?
  2. Hourglass
    • Hourglass最初时用于人体姿态估计
    • Stacked Hourglass把多个漏斗形状的网络级联起来,可以获取多尺度的信息
    1. Residual模块 基础的Residual模块
      • 代码在nets/hourglass/class residual之中。就是简单的残差链接网络
    2. Hourglass子模块
      • kp module指的时hourglass基本模块
      • 代码在nets/hourglass/class kp_model之中
      • 两个主要的函数make_layer(将空间分辨率降维) 和make_layer_revr(将函数升维)
  3. Hourglass exkp时Hourglass模型整体实现
    • 具体代码在nets/hourglass/class exkp之中
    • 注意inters变量,这个变量保存的是中间监督过程,可以在这个位置添加loss。

版权声明:本文为qq_35556504原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。