环境配置

ubuntu系统的参考我这篇博文：点击进入；

一、git clone source code

从官方github仓库上把源码下下来，在pycharm上新建一个名称为“CenterNet”的工程，建好后的目录如下：

二、准备数据集

2-1 在data文件夹下新建两个文件夹：food、image_and_xml；
其中food为你自己数据集的名称，我这里要做的是识别一个菜品的任务，所以命名为food，image_and_xml存放的是你所有的图片和xml文件；
注：这里的xml文件可以下载labelimg标注工具进行标注，网上一堆。

2-2 新建一个python文件，文件代码为xml_to_json，即将voc数据集的xml文件格式转为coco数据集的json文件格式，记得在运行之前要把相应的文件路径以及一些参数设置好；
注意，这里默认划分的是训练、验证和测试集，如果你不需要测试集，只需要简单的改下代码即可，不会改的评论留言下这边指导下即可。
代码如下：

# coding:utf-8
# 运行前请先做以下工作：
# pip install lxml
# 将所有的图片及xml文件存放到xml_dir指定的文件夹下，并将此文件夹放置到当前目录下
#import os
import glob
import json
import shutil
import numpy as np
import xml.etree.ElementTree as ETSTART_BOUNDING_BOX_ID = 1
save_path = "."def get(root, name):return root.findall(name)def get_and_check(root, name, length):vars = get(root, name)if len(vars) == 0:raise NotImplementedError('Can not find %s in %s.' % (name, root.tag))if length and len(vars) != length:raise NotImplementedError('The size of %s is supposed to be %d, but is %d.' % (name, length, len(vars)))if length == 1:vars = vars[0]return varsdef convert(xml_list, json_file):json_dict = {"images": [], "type": "instances", "annotations": [], "categories": []}categories = pre_define_categories.copy()bnd_id = START_BOUNDING_BOX_IDall_categories = {}for index, line in enumerate(xml_list):# print("Processing %s"%(line))xml_f = linetree = ET.parse(xml_f)root = tree.getroot()filename = os.path.basename(xml_f)[:-4] + ".jpg"image_id = 20190000001 + indexsize = get_and_check(root, 'size', 1)width = int(get_and_check(size, 'width', 1).text)height = int(get_and_check(size, 'height', 1).text)image = {'file_name': filename, 'height': height, 'width': width, 'id': image_id}json_dict['images'].append(image)#  Currently we do not support segmentationsegmented = get_and_check(root, 'segmented', 1).textassert segmented == '0'for obj in get(root, 'object'):category = get_and_check(obj, 'name', 1).textif category in all_categories:all_categories[category] += 1else:all_categories[category] = 1if category not in categories:if only_care_pre_define_categories:continuenew_id = len(categories) + 1print("[warning] category '{}' not in 'pre_define_categories'({}), create new id: {} automatically".format(category, pre_define_categories, new_id))categories[category] = new_idcategory_id = categories[category]bndbox = get_and_check(obj, 'bndbox', 1)xmin = int(float(get_and_check(bndbox, 'xmin', 1).text))ymin = int(float(get_and_check(bndbox, 'ymin', 1).text))xmax = int(float(get_and_check(bndbox, 'xmax', 1).text))ymax = int(float(get_and_check(bndbox, 'ymax', 1).text))assert (xmax > xmin), "xmax <= xmin, {}".format(line)assert (ymax > ymin), "ymax <= ymin, {}".format(line)o_width = abs(xmax - xmin)o_height = abs(ymax - ymin)ann = {'area': o_width * o_height, 'iscrowd': 0, 'image_id':image_id, 'bbox': [xmin, ymin, o_width, o_height],'category_id': category_id, 'id': bnd_id, 'ignore': 0,'segmentation': []}json_dict['annotations'].append(ann)bnd_id = bnd_id + 1for cate, cid in categories.items():cat = {'supercategory': 'food', 'id': cid, 'name': cate}json_dict['categories'].append(cat)json_fp = open(json_file, 'w')json_str = json.dumps(json_dict)json_fp.write(json_str)json_fp.close()print("------------create {} done--------------".format(json_file))print("find {} categories: {} -->>> your pre_define_categories {}: {}".format(len(all_categories),all_categories.keys(),len(pre_define_categories),pre_define_categories.keys()))print("category: id --> {}".format(categories))print(categories.keys())print(categories.values())if __name__ == '__main__':# 定义你自己的类别classes = ['aaa', 'bbb', 'ccc', 'ddd', 'eee', 'fff']pre_define_categories = {}for i, cls in enumerate(classes):pre_define_categories[cls] = i + 1# 这里也可以自定义类别id，把上面的注释掉换成下面这行即可# pre_define_categories = {'a1': 1, 'a3': 2, 'a6': 3, 'a9': 4, "a10": 5}only_care_pre_define_categories = True  # or False# 保存的json文件save_json_train = 'train_food.json'save_json_val = 'val_food.json'save_json_test = 'test_food.json'# 初始文件所在的路径xml_dir = "./image_and_xml"xml_list = glob.glob(xml_dir + "/*.xml")xml_list = np.sort(xml_list)# 打乱数据集np.random.seed(100)np.random.shuffle(xml_list)# 按比例划分打乱后的数据集train_ratio = 0.8val_ratio = 0.1train_num = int(len(xml_list) * train_ratio)val_num = int(len(xml_list) * val_ratio)xml_list_train = xml_list[:train_num]xml_list_val = xml_list[train_num: train_num+val_num]xml_list_test = xml_list[train_num+val_num:]# 将xml文件转为coco文件，在指定目录下生成三个json文件（train/test/food）convert(xml_list_train, save_json_train)convert(xml_list_val, save_json_val)convert(xml_list_test, save_json_test)# # 将图片按照划分后的结果进行存放# if os.path.exists(save_path + "/annotations"):#     shutil.rmtree(save_path + "/annotations")# os.makedirs(save_path + "/annotations")# if os.path.exists(save_path + "/images_divide/train"):#     shutil.rmtree(save_path + "/images_divide/train")# os.makedirs(save_path + "/images_divide/train")# if os.path.exists(save_path + "/images_divide/val"):#     shutil.rmtree(save_path + "/images_divide/val")# os.makedirs(save_path + "/images_divide/val")# if os.path.exists(save_path + "/images_divide/test"):#     shutil.rmtree(save_path + "/images_divide/test")# os.makedirs(save_path + "/images_divide/test")# # 按需执行，生成3个txt文件，存放相应的文件名称# f1 = open("./train.txt", "w")# for xml in xml_list_train:#     img = xml[:-4] + ".jpg"#     f1.write(os.path.basename(xml)[:-4] + "\n")#     shutil.copyfile(img, save_path + "/images_divide/train/" + os.path.basename(img))## f2 = open("val.txt", "w")# for xml in xml_list_val:#     img = xml[:-4] + ".jpg"#     f2.write(os.path.basename(xml)[:-4] + "\n")#     shutil.copyfile(img, save_path + "/images_divide/val/" + os.path.basename(img))## f3 = open("test.txt", "w")# for xml in xml_list_val:#     img = xml[:-4] + ".jpg"#     f2.write(os.path.basename(xml)[:-4] + "\n")#     shutil.copyfile(img, save_path + "/images_divide/test/" + os.path.basename(img))## f1.close()# f2.close()# f3.close()print("-" * 50)print("train number:", len(xml_list_train))print("val number:", len(xml_list_val))print("test number:", len(xml_list_val))

运行完之后，会得到三个json文件，分别代表训练，测试和验证。

2-3 进入到food数据集下，新建两个文件夹：images和annotations：
images：存放你的所有图片文件；
annotations：把上一步生成的三个json文件复制或剪切到这个文件夹下；

三、修改配置信息

3-1 计算所有的图片的均值和标准差，直接将图片存放到同一个文件夹，把路径改下即可：

import cv2, os, argparse
import numpy as np
from tqdm import tqdmdef main():dirs = r'F:\Pycharm Professonal\CenterNet\CenterNet\data\food\images'  # 修改你自己的图片路径img_file_names = os.listdir(dirs)m_list, s_list = [], []for img_filename in tqdm(img_file_names):img = cv2.imread(dirs + '/' + img_filename)img = img / 255.0m, s = cv2.meanStdDev(img)m_list.append(m.reshape((3,)))s_list.append(s.reshape((3,)))m_array = np.array(m_list)s_array = np.array(s_list)m = m_array.mean(axis=0, keepdims=True)s = s_array.mean(axis=0, keepdims=True)print("mean = ", m[0][::-1])print("std = ", s[0][::-1])if __name__ == '__main__':main()

mean =  [0.43708543 0.4406526  0.42904118]
std =  [0.25848474 0.25735703 0.2562089 ]

3-2 写一个数据类
到src/lib/datasets/dataset目录下，新建一个python文件，这里需要自己写一个数据类，我这里命名为food.py，；
（1）第14行的类名改为自己的类型名，这里定义为Food；
（2）第15行的num_class改为自己数据集的类别数；
（3）第16行的default_resolution为默认的分辨率，这里按原作者给出的[512, 512]，如果觉得自己的硬件设备跟不上，可以适当的改小，注意上面所计算出来的整个数据集的均值和标准差也要同步；
（4）第17-20行的均值和方差填上去；
（5）第23行super类的继承改为你自己定义的类名称；
（6）修改读取json文件的路径；
（7）修改类别名字和id；
总的可参考下面的代码：

from __future__ import absolute_import
from __future__ import division
from __future__ import print_functionimport pycocotools.coco as coco
from pycocotools.cocoeval import COCOeval
import numpy as np
import json
import osimport torch.utils.data as dataclass Food(data.Dataset):num_classes = 6default_resolution = [512, 512]mean = np.array([0.472459, 0.475080, 0.482652],dtype=np.float32).reshape((1, 1, 3))std = np.array([0.255084, 0.254665, 0.257073],dtype=np.float32).reshape((1, 1, 3))def __init__(self, opt, split):super(Food, self).__init__()self.data_dir = os.path.join(opt.data_dir, 'food')self.img_dir = os.path.join(self.data_dir, 'images')if split == 'val':self.annot_path = os.path.join(self.data_dir, 'annotations', 'val_food.json')else:if opt.task == 'exdet':self.annot_path = os.path.join(self.data_dir, 'annotations', 'train_food.json')if split == 'test':self.annot_path = os.path.join(self.data_dir, 'annotations', 'test_food.json')else:self.annot_path = os.path.join(self.data_dir, 'annotations', 'train_food.json')self.max_objs = 128self.class_name = ['__background__', 'aaa', 'bbb', 'ccc', 'ddd', 'eee', 'fff']self._valid_ids = [1, 2, 3, 4, 5, 6]self.cat_ids = {v: i for i, v in enumerate(self._valid_ids)}self.voc_color = [(v // 32 * 64 + 64, (v // 8) % 4 * 64, v % 8 * 32) \for v in range(1, self.num_classes + 1)]self._data_rng = np.random.RandomState(123)self._eig_val = np.array([0.2141788, 0.01817699, 0.00341571],dtype=np.float32)self._eig_vec = np.array([[-0.58752847, -0.69563484, 0.41340352],[-0.5832747, 0.00994535, -0.81221408],[-0.56089297, 0.71832671, 0.41158938]], dtype=np.float32)# self.mean = np.array([0.485, 0.456, 0.406], np.float32).reshape(1, 1, 3)# self.std = np.array([0.229, 0.224, 0.225], np.float32).reshape(1, 1, 3)self.split = splitself.opt = optprint('==> initializing food {} data.'.format(split))self.coco = coco.COCO(self.annot_path)self.images = self.coco.getImgIds()self.num_samples = len(self.images)print('Loaded {} {} samples'.format(split, self.num_samples))@staticmethoddef _to_float(x):return float("{:.2f}".format(x))def convert_eval_format(self, all_bboxes):# import pdb; pdb.set_trace()detections = []for image_id in all_bboxes:for cls_ind in all_bboxes[image_id]:category_id = self._valid_ids[cls_ind - 1]for bbox in all_bboxes[image_id][cls_ind]:bbox[2] -= bbox[0]bbox[3] -= bbox[1]score = bbox[4]bbox_out = list(map(self._to_float, bbox[0:4]))detection = {"image_id": int(image_id),"category_id": int(category_id),"bbox": bbox_out,"score": float("{:.2f}".format(score))}if len(bbox) > 5:extreme_points = list(map(self._to_float, bbox[5:13]))detection["extreme_points"] = extreme_pointsdetections.append(detection)return detectionsdef __len__(self):return self.num_samplesdef save_results(self, results, save_dir):json.dump(self.convert_eval_format(results),open('{}/results.json'.format(save_dir), 'w'))def run_eval(self, results, save_dir):# result_json = os.path.join(save_dir, "results.json")# detections  = self.convert_eval_format(results)# json.dump(detections, open(result_json, "w"))self.save_results(results, save_dir)coco_dets = self.coco.loadRes('{}/results.json'.format(save_dir))coco_eval = COCOeval(self.coco, coco_dets, "bbox")coco_eval.evaluate()coco_eval.accumulate()coco_eval.summarize()

3-3 将数据集加入src/lib/datasets/dataset_factory里面

from __future__ import absolute_import
from __future__ import division
from __future__ import print_functionfrom .sample.ddd import DddDataset
from .sample.exdet import EXDetDataset
from .sample.ctdet import CTDetDataset
from .sample.multi_pose import MultiPoseDatasetfrom .dataset.coco import COCO
from .dataset.pascal import PascalVOC
from .dataset.kitti import KITTI
from .dataset.coco_hp import COCOHP
from .dataset.food import Fooddataset_factory = {'coco': COCO,'pascal': PascalVOC,'kitti': KITTI,'coco_hp': COCOHP,'food': Food
}_sample_factory = {'exdet': EXDetDataset,'ctdet': CTDetDataset,'ddd': DddDataset,'multi_pose': MultiPoseDataset
}def get_dataset(dataset, task):class Dataset(dataset_factory[dataset], _sample_factory[task]):passreturn Dataset

3-4 在/src/lib/opts.py文件中修改

    self.parser.add_argument('--dataset', default='food',help='food | coco | kitti | coco_hp | pascal')

3-5 修改ctdet任务使用的默认数据集为新添加的数据集，如下（修改分辨率，类别数，均值，标准差，数据集名字）：
第336行，改下ctdet的init初始化信息。

  def init(self, args=''):default_dataset_info = {'ctdet': {'default_resolution': [512, 512], 'num_classes': 6,'mean': [0.472459, 0.475080, 0.482652], 'std': [0.255084, 0.254665, 0.257073],'dataset': 'food'},

3-6 修改src/lib/utils/debugger.py文件(变成自己数据的类别和名字，前后数据集名字一定保持一致）
（1）第45行下方加入两行：

    elif num_classes == 6 or dataset == 'food':self.names = food_class_name

（2）第460行下方加入自己所定义的类别，不包含背景：

food_class_name = ['aaa', 'bbb', 'ccc', 'ddd', 'eee', 'fff']

四、训练数据

在./src/目录下，运行main.py文件，这里food改成你自己要保存的实验结果文件夹名称即可：
4.1.1 不加载预训练权重：

python main.py ctdet --exp_id food --batch_size 32 --lr 1.25e-4  --gpus 0

4.1.2 加载预训练权重：

python main.py ctdet --exp_id food --batch_size 32 --lr 1.25e-4  --gpus 0 --load_model ../models/ctdet_dla_2x.pth

4.1.3 多卡训练，其中master_batch_sizes 表示的是你在主GPU上要放置多大的batch_size，其余分配到其它卡上：

python main.py ctdet --exp_id food --batch_size 32 --lr 1.25e-4  --gpus 0,1 --load_model ../models/ctdet_dla_2x.pth --master_batch_size 8

4.1.4 断点恢复训练，比如你将结果保存在这个exp_id=food，那么food文件夹下就会有model_last.pth这个，想继续恢复训练：

python main.py ctdet --exp_id food --batch_size 32 --lr 1.25e-4  --gpus 0 --resume

若提示报错，则尝试修改opts文件中的num_workers改为0，或者将batch_size调小。
训练完成后，在./exp/ctdet/food/文件夹下会出现一堆文件；
其中，model_last是最后一次epoch的模型；model_best是val最好的模型；

五、验证模型

运行demo文件检查下训练的模型，–demo设置你要预测的图片/图片文件夹/视频所在的路径；

5.1 原始预测

python demo.py ctdet --demo ../data/food/images/food1.jpg  --load_model ../exp/ctdet/food/model_best.pth

5.2 带数据增强预测

python demo.py ctdet --demo ../data/food/images/food1.jpg  --load_model ../exp/ctdet/food/model_best.pth --flip_test

5.3 多尺度预测

python demo.py ctdet --demo ../data/food/images/food1.jpg  --load_model ../exp/ctdet/food/model_best.pth --test_scales 0.5,0.75,1.0,1.25,1.5

注意，如果多尺度预测报错，一般就是你自己没有编译nms。
编译方法：到 path/to/CenterNet/src/lib/externels 目录下，运行：

python setup.py build_ext --inplace

如果需要保存你的预测结果，可以到目录 path/to/CenterNet/src/lib/detecors/ctdet.py下，在show_results函数中的末尾加入这句：

# path替换成你所需要保存的路径，并确定这个文件夹是否存在
debugger.save_all_imgs(path='/CenterNet-master/outputs', genID=True)

六、测试数据

python test.py --exp_id food --not_prefetch_test ctdet --load_model ../CenterNet/exp/ctdet/food/model_best.pth

七、批量保存每张图片的预测的结果（bbox，id，score）

7.1 进入到CenterNet/src/lib/utils/debugger.py，Ctrl+F找到add_coco_bbox()这个方法，将方法替换为：

    def add_coco_bbox(self, bbox, cat, conf=1, show_txt=True, img_id='default'):bbox = np.array(bbox, dtype=np.int32)# cat = (int(cat) + 1) % 80cat = int(cat)# print('cat', cat, self.names[cat])c = self.colors[cat][0][0].tolist()if self.theme == 'white':c = (255 - np.array(c)).tolist()txt = '{}{:.1f}'.format(cat, conf)bbox_info = [int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])]info = [bbox_info, self.names[cat], float(conf)]font = cv2.FONT_HERSHEY_SIMPLEXcat_size = cv2.getTextSize(txt, font, 0.5, 2)[0]cv2.rectangle(self.imgs[img_id], (bbox[0], bbox[1]), (bbox[2], bbox[3]), c, 2)if show_txt:cv2.rectangle(self.imgs[img_id],(bbox[0], bbox[1] - cat_size[1] - 2),(bbox[0] + cat_size[0], bbox[1] - 2), c, -1)cv2.putText(self.imgs[img_id], txt, (bbox[0], bbox[1] - 2),font, 0.5, (0, 0, 0), thickness=1, lineType=cv2.LINE_AA)return info

这里info便保存了每张图片的每个预测框，对应的类别和置信度信息。

7.2- 进入到CenterNet/src/lib/detectors/ctdet.py，这个文件夹当中，找到show_results这个方法，替换为：

    def show_results(self, debugger, image, results):debugger.add_img(image, img_id='ctdet')infos = []for j in range(1, self.num_classes + 1):for bbox in results[j]:if bbox[4] > self.opt.vis_thresh:info = debugger.add_coco_bbox(bbox[:4], j - 1, bbox[4], img_id='ctdet')infos.append(info)debugger.show_all_imgs(pause=self.pause)return infos

7.3 进入到CenterNet/src/lib/detectors/base_detector.py，找到run这个方法，替换为：

def run(self, image_or_path_or_tensor, meta=None):load_time, pre_time, net_time, dec_time, post_time = 0, 0, 0, 0, 0merge_time, tot_time = 0, 0debugger = Debugger(dataset=self.opt.dataset, ipynb=(self.opt.debug == 3),theme=self.opt.debugger_theme)start_time = time.time()pre_processed = Falseif isinstance(image_or_path_or_tensor, np.ndarray):image = image_or_path_or_tensorelif type(image_or_path_or_tensor) == type(''):image = cv2.imread(image_or_path_or_tensor)else:image = image_or_path_or_tensor['image'][0].numpy()pre_processed_images = image_or_path_or_tensorpre_processed = Trueloaded_time = time.time()load_time += (loaded_time - start_time)detections = []for scale in self.scales:  # scales = [1]scale_start_time = time.time()if not pre_processed:# 运行这里images, meta = self.pre_process(image, scale, meta)else:# import pdb; pdb.set_trace()images = pre_processed_images['images'][scale][0]meta = pre_processed_images['meta'][scale]meta = {k: v.numpy()[0] for k, v in meta.items()}images = images.to(self.opt.device)torch.cuda.synchronize()pre_process_time = time.time()pre_time += pre_process_time - scale_start_timeoutput, dets, forward_time = self.process(images, return_time=True)torch.cuda.synchronize()net_time += forward_time - pre_process_timedecode_time = time.time()dec_time += decode_time - forward_timeif self.opt.debug >= 2:self.debug(debugger, images, dets, output, scale)dets = self.post_process(dets, meta, scale)torch.cuda.synchronize()post_process_time = time.time()post_time += post_process_time - decode_timedetections.append(dets)results = self.merge_outputs(detections)torch.cuda.synchronize()end_time = time.time()merge_time += end_time - post_process_timetot_time += end_time - start_timeif self.opt.debug >= 1:info = self.show_results(debugger, image, results)return {'results': results, 'tot': tot_time, 'load': load_time,'pre': pre_time, 'net': net_time, 'dec': dec_time,'post': post_time, 'merge': merge_time}, info

7.4 最后将CenterNet/src/demo.py 这个文件的内容替换为：

from __future__ import absolute_import
from __future__ import division
from __future__ import print_functionimport _init_pathsimport os
import cv2
import jsonfrom opts import opts
from detectors.detector_factory import detector_factoryimage_ext = ['jpg', 'jpeg', 'png', 'webp']
video_ext = ['mp4', 'mov', 'avi', 'mkv']
time_stats = ['tot', 'load', 'pre', 'net', 'dec', 'post', 'merge']def demo():os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_stropt.debug = max(opt.debug, 1)Detector = detector_factory[opt.task]detector = Detector(opt)if opt.demo == 'webcam' or \opt.demo[opt.demo.rfind('.') + 1:].lower() in video_ext:cam = cv2.VideoCapture(0 if opt.demo == 'webcam' else opt.demo)detector.pause = Falsewhile True:_, img = cam.read()cv2.imshow('input', img)ret = detector.run(img)time_str = ''for stat in time_stats:time_str = time_str + '{} {:.3f}s |'.format(stat, ret[stat])print(time_str)if cv2.waitKey(1) == 27:return  # esc to quitelse:if os.path.isdir(opt.demo):image_names = []ls = os.listdir(opt.demo)for file_name in sorted(ls):ext = file_name[file_name.rfind('.') + 1:].lower()if ext in image_ext:image_names.append(os.path.join(opt.demo, file_name))else:image_names = [opt.demo]results = {}for (image_name) in image_names:ret, info = detector.run(image_name)save_name = image_name.split('/')[-1]results[save_name] = infotime_str = ''for stat in time_stats:time_str = time_str + '{} {:.3f}s |'.format(stat, ret[stat])print(time_str)results_str = json.dumps(results)with open(opt.save_dir+"/{}.json".format(opt.exp_id), 'w') as json_file:json_file.write(results_str)if __name__ == '__main__':opt = opts().init()demo()

上面将预测信息保存为一个json文件，保存路径可自己设置。

知乎ID：https://www.zhihu.com/people/peissen，欢迎关注讨论。