CenterNet2实战：手把手带你实现使用CenterNet2训练自定义数据集

小灰灰 2022-09-13 11:27 353阅读 0赞

# 1、CenterNet2 介绍 #

论文地址：[https://arxiv.org/abs/2103.07461 ][https_arxiv.org_abs_2103.07461]

GitHub地址：https://github.com/xingyizhou/CenterNet2

CenterNet的原作者提出了一个概率性的两阶段检测器， 这种解释激发了强大的第一阶段的使用，该阶段学习估计对象可能性而不是最大化召回。 然后将这些可能性与来自第二阶段的分类分数相结合，以产生最终检测的原则概率分数。 概率两级检测器比一级或两级检测器更快、更准确，作者将其命名为CenterNet2。CenterNet2包含多个模型，COCO成绩最高达到56.4mAP，超越了ScaledYOLOV4，成为现阶段最强的物体检测模型。

![image-20210927170129602][]

这个模型是在Detectron2上开发的，版本是2.3。如果是初次使用需要了解一下Detectron2。

# 2、Detectron2介绍 #

**Detectron2** 前身就是鼎鼎大名的 **Detectron**，其实Detectron可以说是Facebook第一代检测工具箱。Detectron2 不仅支持 Detectron已有的目标检测、实例分割、姿态估计等任务，还支持语义分割和全景分割。

优点如下：

**基于PyTorch**：PyTorch可以提供更直观的命令式编程模型，开发者可以更快的进行迭代模型设计和实验。

**模块化、可扩展**：从Detectron2开始，Facebook引入了模块化设计，允许用户将自定义模块插入目标检测系统的几乎任何部分。这意味着许多新的研究项目和核心Detectron2库可以完全分开。其可扩展性也使得Detectron2更加灵活。

**支持语义分割和全景分割**。

**实现质量**：从头开始重写推出的Detectron2，使得能够重新审视低级设计决策并解决了原始Detectron中的几个实现问题。

**速度和可扩展性**：Detectron2比原始Detectron更快，而且可以更加方便进行GPU服务器的分布式训练。

**Detectron2go**：新增了将模型产品化部署的软件实现，包括标准的内部数据训练工作流实现、模型压缩量化、模型转化等。

总之，我们使用Detectron2很方便的实现模型的训练、测试以及模型转换。所以现在很多的新模型都是在Detectron2开发。

# 3、搭建CenterNet2 测试环境 #

我本地环境：

操作系统：win10、Cuda11.0。

## 3.1 创建虚拟环境 ##

创建虚拟环境，并激活环境。

conda create --name centernet2 python=3.7
    activate centernet2
    conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=11.0 -c pytorch

## 3.2 安装apex ##

APEX是英伟达开源的，完美支持PyTorch框架，用于改变数据格式来减小模型显存占用的工具。其中最有价值的是amp（Automatic Mixed Precision），将模型的大部分操作都用Float16数据类型测试，一些特别操作仍然使用Float32。并且用户仅仅通过三行代码即可完美将自己的训练代码迁移到该模型。实验证明，使用Float16作为大部分操作的数据类型，并没有降低参数，在一些实验中，反而由于可以增大Batch size，带来精度上的提升，以及训练速度上的提升。

### 3.2.1 下载apex ###

网址 https://github.com/NVIDIA/apex，下载到本地文件夹。解压后进入到apex的目录安装依赖。在执行命令；

cd C:\Users\WH\Downloads\apex-master #进入apex目录
    pip install -r requirements.txt

### 3.2.2 安装apex ###

依赖安装完后，打开cmd，cd进入到刚刚下载完的apex-master路径下，运行：

python setup.py install

然后跑了一堆东西，最后是这样的：  
![image-20210928065758887][]

安装完成！

## 3.3 安装fvcore ##

fvcore库的简介  
fvcore是一个轻量级的核心库，它提供了在各种计算机视觉框架(如Detectron2)中共享的最常见和最基本的功能。这个库基于Python 3.6+和PyTorch。这个库中的所有组件都经过了类型注释、测试和基准测试。Facebook 的人工智能实验室即FAIR的计算机视觉组负责维护这个库。

github地址：https://github.com/facebookresearch/fvcore

执行命令

conda install -c fvcore -c iopath -c conda-forge fvcore

## 3.4 安装其他的库 ##

安装pycocotools

pip install pycocotools

安装cv2

pip install opencv-python

安装 antlr4

pip install antlr4-python3-runtime

安装future

pip install future

安装protobuf

pip install protobuf

安装absl

pip install absl-py

3.5 编译CenterNet2

进入CenterNet2目录，目录根据自己的实际情况更改

cd D:\CenterNet2-master

编译

python setup.py install

# 4、测试环境 #

新建imgs和imgout文件夹，imgs文件夹存放待测试的图片。

图片如下：

\[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-ymnqnmPL-1632828249507)(https://gitee.com/wanghao1090220084/images/raw/master/img/image-20210928095744639.png)\]

下载模型“CenterNet2\_R50\_1x.pth”(注：作者将模型放到谷歌的网盘上，这个国内访问不了，如果不能下载，关注我的公众号获取连接。)，将其放到“projects/CenterNet2”目录下面。

执行命令：

python projects/CenterNet2/demo.py --config-file projects/CenterNet2/configs/CenterNet2_R50_1x.yaml --input imgs/ --output imgout --opts MODEL.WEIGHTS projects/CenterNet2/CenterNet2_R50_1x.pth

运行结果：

![image-20210928095902157][]

能够运行demo说明环境已经没有问题了。

# 5、制作数据集 #

本次采用的数据集是Labelme标注的数据集，地址：链接：https://pan.baidu.com/s/1nxo9-NpNWKK4PwDZqwKxGQ 提取码：kp4e，需要将其转为COCO格式的数据集。转换代码如下：

新建labelme2coco.py

import argparse
    import json
    import matplotlib.pyplot as plt
    import skimage.io as io
    import cv2
    from labelme import utils
    import numpy as np
    import glob
    import PIL.Image
    
    REQUIRE_MASK = False
    labels = { 'aircraft': 1, 'oiltank': 2}
    
    
    class labelme2coco(object):
        def __init__(self, labelme_json=[], save_json_path='./new.json'):
            ''' :param labelme_json: the list of all labelme json file paths :param save_json_path: the path to save new json '''
            self.labelme_json = labelme_json
            self.save_json_path = save_json_path
            self.images = []
            self.categories = []
            self.annotations = []
            # self.data_coco = {}
            self.label = []
            self.annID = 1
            self.height = 0
            self.width = 0
            self.require_mask = REQUIRE_MASK
            self.save_json()
    
        def data_transfer(self):
            for num, json_file in enumerate(self.labelme_json):
                if not json_file == self.save_json_path:
                    with open(json_file, 'r') as fp:
                        data = json.load(fp)
                        self.images.append(self.image(data, num))
                        for shapes in data['shapes']:
                            print("label is ")
                            print(shapes['label'])
                            label = shapes['label']
                            # if label[1] not in self.label:
                            if label not in self.label:
                                print("find new category: ")
                                self.categories.append(self.categorie(label))
                                print(self.categories)
                                # self.label.append(label[1])
                                self.label.append(label)
                            points = shapes['points']
                            self.annotations.append(self.annotation(points, label, num))
                            self.annID += 1
    
        def image(self, data, num):
            image = { }
            img = utils.img_b64_to_arr(data['imageData'])
            height, width = img.shape[:2]
            img = None
            image['height'] = height
            image['width'] = width
            image['id'] = num + 1
            image['file_name'] = data['imagePath'].split('/')[-1]
            self.height = height
            self.width = width
            return image
    
        def categorie(self, label):
            categorie = { }
            categorie['supercategory'] = label
            # categorie['supercategory'] = label
            categorie['id'] = labels[label]  # 0 默认为背景
            categorie['name'] = label
            return categorie
    
        def annotation(self, points, label, num):
            annotation = { }
            print(points)
            x1 = points[0][0]
            y1 = points[0][1]
            x2 = points[1][0]
            y2 = points[1][1]
            contour = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]])  # points = [[x1, y1], [x2, y2]] for rectangle
            contour = contour.astype(int)
            area = cv2.contourArea(contour)
            print("contour is ", contour, " area = ", area)
            annotation['segmentation'] = [list(np.asarray([[x1, y1], [x2, y1], [x2, y2], [x1, y2]]).flatten())]
            # [list(np.asarray(contour).flatten())]
            annotation['iscrowd'] = 0
            annotation['area'] = area
            annotation['image_id'] = num + 1
    
            if self.require_mask:
                annotation['bbox'] = list(map(float, self.getbbox(points)))
            else:
                x1 = points[0][0]
                y1 = points[0][1]
                width = points[1][0] - x1
                height = points[1][1] - y1
                annotation['bbox'] = list(np.asarray([x1, y1, width, height]).flatten())
    
            annotation['category_id'] = self.getcatid(label)
            annotation['id'] = self.annID
            return annotation
    
        def getcatid(self, label):
            for categorie in self.categories:
                # if label[1]==categorie['name']:
                if label == categorie['name']:
                    return categorie['id']
            return -1
    
        def getbbox(self, points):
            polygons = points
            mask = self.polygons_to_mask([self.height, self.width], polygons)
            return self.mask2box(mask)
    
        def mask2box(self, mask):
    
            # np.where(mask==1)
            index = np.argwhere(mask == 1)
            rows = index[:, 0]
            clos = index[:, 1]
    
            left_top_r = np.min(rows)  # y
            left_top_c = np.min(clos)  # x
            right_bottom_r = np.max(rows)
            right_bottom_c = np.max(clos)
            return [left_top_c, left_top_r, right_bottom_c - left_top_c, right_bottom_r - left_top_r]
    
        def polygons_to_mask(self, img_shape, polygons):
            mask = np.zeros(img_shape, dtype=np.uint8)
            mask = PIL.Image.fromarray(mask)
            xy = list(map(tuple, polygons))
            PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
            mask = np.array(mask, dtype=bool)
            return mask
    
        def data2coco(self):
            data_coco = { }
            data_coco['images'] = self.images
            data_coco['categories'] = self.categories
            data_coco['annotations'] = self.annotations
            return data_coco
    
        def save_json(self):
            print("in save_json")
            self.data_transfer()
            self.data_coco = self.data2coco()
    
            print(self.save_json_path)
            json.dump(self.data_coco, open(self.save_json_path, 'w'), indent=4)
    
    
    labelme_json = glob.glob('LabelmeData/*.json')
    from sklearn.model_selection import train_test_split
    
    trainval_files, test_files = train_test_split(labelme_json, test_size=0.2, random_state=55)
    import os
    
    if not os.path.exists("projects/CenterNet2/datasets/coco/annotations"):
        os.makedirs("projects/CenterNet2/datasets/coco/annotations/")
    if not os.path.exists("projects/CenterNet2/datasets/coco/train2017"):
        os.makedirs("projects/CenterNet2/datasets/coco/train2017")
    if not os.path.exists("projects/CenterNet2/datasets/coco/val2017"):
        os.makedirs("projects/CenterNet2/datasets/coco/val2017")
    labelme2coco(trainval_files, 'projects/CenterNet2/datasets/coco/annotations/instances_train2017.json')
    labelme2coco(test_files, 'projects/CenterNet2/datasets/coco/annotations/instances_val2017.json')
    import shutil
    
    for file in trainval_files:
        shutil.copy(os.path.splitext(file)[0] + ".jpg", "projects/CenterNet2/datasets/coco/train2017/")
    for file in test_files:
        shutil.copy(os.path.splitext(file)[0] + ".jpg", "projects/CenterNet2/datasets/coco/val2017/")

# 6、配置训练环境 #

## 6.1 更改预训练模型的size ##

在projects/CenterNet2目录，新建change\_model\_size.py文件

import torch
    import numpy as np
    import pickle
    num_class = 2
    pretrained_weights  = torch.load('models/CenterNet2_R50_1x.pth')
    pretrained_weights['iteration']=0
    pretrained_weights['model']["roi_heads.box_predictor.0.cls_score.weight"].resize_(num_class+1,1024)
    pretrained_weights['model']["roi_heads.box_predictor.0.cls_score.bias"].resize_(num_class+1)
    pretrained_weights['model']["roi_heads.box_predictor.1.cls_score.weight"].resize_(num_class+1,1024)
    pretrained_weights['model']["roi_heads.box_predictor.1.cls_score.bias"].resize_(num_class+1)
    pretrained_weights['model']["roi_heads.box_predictor.2.cls_score.weight"].resize_(num_class+1,1024)
    pretrained_weights['model']["roi_heads.box_predictor.2.cls_score.bias"].resize_(num_class+1)
    torch.save(pretrained_weights, "models/CenterNet2_%d.pth"%num_class)

这个文件的目的是修改模型输出的size，numclass按照本次打算训练的数据集的类别设置。

## 6.2 修改config参数 ##

路径：“detectron2/engine/defaults.py”

–config-file：模型的配置文件，CenterNet2的模型配置文件放在“projects/CenterNet2/configs”下面。名字和预训练模型对应。

parser.add_argument("--config-file", default="./configs/CenterNet2_DLA-BiFPN-P3_4x.yaml", metavar="FILE", help="path to config file")

resume 是否再次，训练，如果设置为true，则接着上次训练的结果训练。所以第一次训练不用设置。

parser.add_argument(
        "--resume",
        action="store_true",
        help="Whether to attempt to resume from the checkpoint directory. "
        "See documentation of `DefaultTrainer.resume_or_load()` for what it means.",
    )

–num-gpus，gpu的个数，如果只有一个设置为1，如果有多个，可以自己设置想用的个数。

parser.add_argument("--num-gpus", type=int, default=1, help="number of gpus *per machine*")

opts指的是yaml文件的参数。

上面的参数可以设置，也可以不设置，设置之后可以直接运行不用再考虑设置参数，如果不设置每次训练的时候配置一次参数。

修改类别，文件路径“projects/CenterNet2/centernet/config.py”，

_C.MODEL.CENTERNET.NUM_CLASSES = 2

![image-20210928125813104][]

修改yaml文件参数

Base-CenterNet2.yaml中修改预训练模型的路径。

WEIGHTS: "CenterNet2_2.pth"

BASE\_LR：设置学习率。

STEPS：设置训练多少步之后调整学习率。

MAX\_ITER：最大迭代次数。

CHECKPOINT\_PERIOD：设置迭代多少次保存一次模型

BASE_LR: 0.01
    STEPS: (10000, 50000)
    MAX_ITER: 100000
    CHECKPOINT_PERIOD: 5000

在设置上面的参数时要注意，如果选择用CenterNet2\_R50\_1x.yaml，里面没有参数，则在Base-CenterNet2.yaml中设置，如果选用其他的，例如CenterNet2\_DLA-BiFPN-P3\_4x.yaml，这些参数需要在CenterNet2\_DLA-BiFPN-P3\_4x.yaml改。

## 6.3 修改train\_net.py ##

主要修改该setup函数，增加数据集注册。

NUM_CLASSES=2
    
    def setup(args):
        """ Create configs and perform basic setups. """
        register_coco_instances("train", { }, "datasets/coco/annotations/instances_train2017.json",
                                "datasets/coco/train2017")
        register_coco_instances("test", { }, "datasets/coco/annotations/instances_val2017.json",
                                "datasets/coco/val2017")
    
    
        cfg = get_cfg()
        add_centernet_config(cfg)
        cfg.merge_from_file(args.config_file)
        cfg.merge_from_list(args.opts)
        cfg.DATASETS.TRAIN = ("train",)
        cfg.DATASETS.TEST = ("test",)
        cfg.MODEL.CENTERNET.NUM_CLASSES = NUM_CLASSES
        cfg.MODEL.ROI_HEADS.NUM_CLASSES = NUM_CLASSES
    
    
        if '/auto' in cfg.OUTPUT_DIR:
            file_name = os.path.basename(args.config_file)[:-5]
            cfg.OUTPUT_DIR = cfg.OUTPUT_DIR.replace('/auto', '/{}'.format(file_name))
            logger.info('OUTPUT_DIR: {}'.format(cfg.OUTPUT_DIR))
        cfg.freeze()
        default_setup(cfg, args)
        return cfg

还要修改detectron2/engine/launch.py，在launch函数下面增加一句

dist.init_process_group('gloo', init_method='file://tmp/somefile', rank=0, world_size=1)

如下图：

![image-20210928131935858][]

这句话的作用是初始化分布式训练，因为我们没有使用分布式，所以没有初始化，但是不初始化就会报错，所以加上这句。

# 7、训练 #

两种启动方式：

第一种，命令行：进入“projects/CenterNet2/”目录下,执行：

python train_net.py

第二种，直接在pycharm 直接运行train\_net.py.

训练结果：

![image-20210927083348009][]

从训练结果上看，效果确实不错，不过模型很大。大约有500M

![image-20210928132823043][]

# 8、测试 #

修改projects/CenterNet2/demo.py

## 8.1 修改setup\_cfg函数 ##

![image-20210928140145057][]

在红框的位置增加代码，详细如下面的代码。

NUM_CLASSES=2
    def setup_cfg(args):
        # load config from file and command-line arguments
        cfg = get_cfg()
        add_centernet_config(cfg)
        cfg.MODEL.CENTERNET.NUM_CLASSES = NUM_CLASSES
        cfg.MODEL.ROI_HEADS.NUM_CLASSES = NUM_CLASSES
    
        cfg.merge_from_file(args.config_file)
        cfg.merge_from_list(args.opts)
        # Set score_threshold for builtin models
        cfg.MODEL.RETINANET.SCORE_THRESH_TEST = args.confidence_threshold
        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = args.confidence_threshold
        if cfg.MODEL.META_ARCHITECTURE in ['ProposalNetwork', 'CenterNetDetector']:
            cfg.MODEL.CENTERNET.INFERENCE_TH = args.confidence_threshold
            cfg.MODEL.CENTERNET.NMS_TH = cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST
        cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = args.confidence_threshold
        cfg.freeze()
        return cfg

## 8.2 修改显示类别 ##

![image-20210928141246325][]

代码：

visualizer.metadata.thing_classes[:10] = ["aircraft", "oiltank"]

然后进入CenterNet2-master目录，执行如下命令：

python projects/CenterNet2/demo.py --config-file projects/CenterNet2/configs/CenterNet2_R50_1x.yaml --input imgs/ --output imgout --opts MODEL.WEIGHTS projects/CenterNet2/output/CenterNet2/CenterNet2_R50_1x/model_final.pth

运行结果：

![image-20210928141739328][]  
关注公众号“AI小浩”，回复“centernet2实战”，获取本次项目源码和本文的PDF版本。

[https_arxiv.org_abs_2103.07461]: https://arxiv.org/abs/2103.07461
[image-20210927170129602]: /images/20220828/eceb4ae431b9493199697d9409edba8c.png
[image-20210928065758887]: /images/20220828/f01934f216344cae97d2d8e76e39dcd3.png
[image-20210928095902157]: /images/20220828/6b6eda0a0a2d4f3e94d003d4e6697e7c.png
[image-20210928125813104]: /images/20220828/4e1f8ea29ec9460faeaafc199c839dd8.png
[image-20210928131935858]: /images/20220828/a41a1a91969c494495e1cd2a43a41e28.png
[image-20210927083348009]: /images/20220828/1d15370e8c8c42cf9c8400c2b6286d74.png
[image-20210928132823043]: /images/20220828/3368e467988a4f1085af14bc7ce436ee.png
[image-20210928140145057]: /images/20220828/a014cd7c85784936b82b9736b78718a5.png
[image-20210928141246325]: /images/20220828/e07d6ab7418d49b2bd74c90dc6ce7504.png
[image-20210928141739328]: /images/20220828/39cc12a0c23e42eb91368f8281afe84e.png