【详细】CenterNet源码分析(Objects as Points)

CenterNet是一种Anchor Free的目标检测方法，其对应论文为《Objects as Points》，论文链接点击进入，本文主要针对CenterNet源码进行分析。具体的算法过程和代码的环境配置过程可以借鉴其他博客内容。

源码链接点击进入

源码中的模型存储在Google云盘里，不是很好下载，我搜集了一些预训练模型，网盘链接为链接：

https://pan.baidu.com/s/1YBbJYbObcSKIRUr5U_0lyg
提取码：lbnf
（包含Res101、18 和 dla 34），如果有其他朋友下载的模型，欢迎评论区给出

源码结构

其中 src 为源码文件夹，experiments 为一些模型训练测试的语句(*.sh)，主要关注这两个文件夹就行，但是源码主要在 src 中。下面就从 src 文件夹一步步来说。

demo.py 主要在配置完环境后，用来测试您的环境是否安装正确。
main.py 主要是进行训练模型
test.py 主要用来进行测试模型效果
lib 文件夹为 CenterNet 的主要代码，包含了数据预处理，模型，后处理和评价

tools文件夹主要是针对数据集的一些工具类。

从main.py谈起

1 加载一些配置

  torch.manual_seed(opt.seed) #设置随机种子
  torch.backends.cudnn.benchmark = not opt.not_cuda_benchmark and not opt.test#能否使用 cudnn 进行加速
  Dataset = get_dataset(opt.dataset, opt.task)#数据集配置
  opt = opts().update_dataset_info_and_set_heads(opt, Dataset)#加载配置信息
  print(opt)
  logger = Logger(opt)#生成配置信息的logger日志
  os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str#0,1,2代表GPU的索引号
  opt.device = torch.device('cuda' if opt.gpus[0] >= 0 else 'cpu')#是否有可用GPU

2 加载模型

  print('Creating model...')
  #创建模型 arch：网络结构名称 heads：网络的头部 head_conv:头部输出的通道个数
  model = create_model(opt.arch, opt.heads, opt.head_conv)
  # 使用 Adam优化器
  optimizer = torch.optim.Adam(model.parameters(), opt.lr)
  #设置初始轮次
  start_epoch = 0
  #如果加载的模型不为空，则家在模型
  if opt.load_model != '':
    model, optimizer, start_epoch = load_model(
      model, opt.load_model, optimizer, opt.resume, opt.lr, opt.lr_step)
  #得到对应任务的训练器
  Trainer = train_factory[opt.task]
  trainer = Trainer(opt, model, optimizer)
  #将模型加载到GPU上
  trainer.set_device(opt.gpus, opt.chunk_sizes, opt.device)

以上关键的代码是create_model和load_model，以及Trainer，这里对这三个进行分别展开剖析

2.1 create_model

位于 lib/models/model.py中

#模型工厂
_model_factory = {
  'res': get_pose_net, # default Resnet with deconv
  'dlav0': get_dlav0, # default DLAup
  'dla': get_dla_dcn,
  'resdcn': get_pose_net_dcn,
  'hourglass': get_large_hourglass_net,
}
# 创建模型
def create_model(arch, heads, head_conv):
  num_layers = int(arch[arch.find('_') + 1:]) if '_' in arch else 0
  arch = arch[:arch.find('_')] if '_' in arch else arch
  get_model = _model_factory[arch]
  model = get_model(num_layers=num_layers, heads=heads, head_conv=head_conv)
  return model

以上关键在 get_model 方法，这里以backbone 为 dla34 为例进行展开，其 get_model 即为下面的代码段，即加载 DLASeg 这个方法：

def get_pose_net(num_layers, heads, head_conv=256, down_ratio=4):
  model = DLASeg('dla{}'.format(num_layers), heads,
                 pretrained=True,
                 down_ratio=down_ratio,
                 final_kernel=1,
                 last_level=5,
                 head_conv=head_conv)
  return model

整体上的网络结构是论文中的附件中的(d)这个插图，具体的代码结构，可以看相关代码

2.2 load_model

def load_model(model, model_path, optimizer=None, resume=False, 
               lr=None, lr_step=None):
  start_epoch = 0
  checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)# 加载模型，checkpoint是个字典
  print('loaded {}, epoch {}'.format(model_path, checkpoint['epoch']))
  #数据字典
  state_dict_ = checkpoint['state_dict']
  state_dict = {}
  #更改 checkpoint 的键名
  # convert data_parallal to model
  for k in state_dict_:
    if k.startswith('module') and not k.startswith('module_list'):
      state_dict[k[7:]] = state_dict_[k]
    else:
      state_dict[k] = state_dict_[k]
  #得到模型的状态字典
  model_state_dict = model.state_dict()
  # check loaded parameters and created model parameters
  msg = 'If you see this, your model does not fully load the ' + \
        'pre-trained weight. Please make sure ' + \
        'you have correctly specified --arch xxx ' + \
        'or set the correct --num_classes for your own dataset.'
  for k in state_dict:# 遍历预训练模型的参数字典
    if k in model_state_dict:#如果在模型的字典中
       #判断 shape 是否一致，如果不一致 将模型中的参数赋值给state_dict中对应的值
      if state_dict[k].shape != model_state_dict[k].shape:
        print('Skip loading parameter {}, required shape{}, '\
              'loaded shape{}. {}'.format(
          k, model_state_dict[k].shape, state_dict[k].shape, msg))
        state_dict[k] = model_state_dict[k]
    else:
      #如果字段不在模型中，则不进行加载
      print('Drop parameter {}.'.format(k) + msg)
  #遍历模型的字段
  for k in model_state_dict:
    #对于不在模型中的字段，将模型初始化的值添加到预训练的模型中
    if not (k in state_dict):
      print('No param {}.'.format(k) + msg)
      state_dict[k] = model_state_dict[k]
  # 从state_dict中将参数加载到模型，如果存在模型和state_dict中不对应的参数，不会报错，因为strict=False
  model.load_state_dict(state_dict, strict=False)
  ### 以上部分从预训练模型加载参数到模型中，其实为我们后来更改模型提供了便利，即不需要严格的将模型与预训练参数一致，在添加新的模块时，只需将参数进行初始化即可

  # resume optimizer parameters
  #当模型训练突然中止时，进行resume模型
  if optimizer is not None and resume:
    if 'optimizer' in checkpoint:
      optimizer.load_state_dict(checkpoint['optimizer'])
      start_epoch = checkpoint['epoch']
      start_lr = lr
      #如果 step 对应学习率更改的 lr_step 则让lr变为原来的1/10
      for step in lr_step:
        if start_epoch >= step:
          start_lr *= 0.1
      #修改优化器中的 学习率
      for param_group in optimizer.param_groups:
        param_group['lr'] = start_lr
      print('Resumed optimizer with start lr', start_lr)
    else:
      print('No optimizer parameters in checkpoint.')
  if optimizer is not None:
    return model, optimizer, start_epoch
  else:
    return model

2.3 Trainer

train_factory中包含了四项任务，这里以检测任务 ctdet为例展开，其他任务可以进行类比

train_factory = {
  'exdet': ExdetTrainer, 
  'ddd': DddTrainer,
  'ctdet': CtdetTrainer,
  'multi_pose': MultiPoseTrainer, 
}
#检测训练器
class CtdetTrainer(BaseTrainer):
  # opt:配置文件 model:模型文件 optimizer：优化器 
  def __init__(self, opt, model, optimizer=None):
    super(CtdetTrainer, self).__init__(opt, model, optimizer=optimizer)
  # 得到损失
  #loss：总损失 hm_loss：热图损失 wh_loss：size损失 off_loss:偏移量损失
  def _get_losses(self, opt):
    loss_states = ['loss', 'hm_loss', 'wh_loss', 'off_loss']
    loss = CtdetLoss(opt)#调用 CtdetLoss类
    return loss_states, loss
  # 执行训练
  def debug(self, batch, output, iter_id):
    opt = self.opt#配置信息
    #是否包含偏移回归值
    reg = output['reg'] if opt.reg_offset else None
    #对训练结果进行整合
    dets = ctdet_decode(
      output['hm'], output['wh'], reg=reg,
      cat_spec_wh=opt.cat_spec_wh, K=opt.K)
    #将检测结果转为cpu下的numpy,并对其进行reshape
    dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])
    dets[:, :, :4] *= opt.down_ratio#对bbox乘以4 让其返回原图中的坐标 
    #加载gt信息
    dets_gt = batch['meta']['gt_det'].numpy().reshape(1, -1, dets.shape[2])
    #对gt的 bbox 也映射回原图
    dets_gt[:, :, :4] *= opt.down_ratio
    #进行一次循环，这个参数可以调整为您的batch_size大小，但是模型速度会变慢，因为需要写结果
    for i in range(1):
    # dataset 数据集 ipynb显示结果的方式 theme：主题风格 默认是白色
      debugger = Debugger(
        dataset=opt.dataset, ipynb=(opt.debug==3), theme=opt.debugger_theme)
      # batch 输入的图像转为 HWC 的形式
      img = batch['input'][i].detach().cpu().numpy().transpose(1, 2, 0)
      # 对输入图像进行反归一化，由于在送入网络时，进行了一次归一化
      img = np.clip(((
        img * opt.std + opt.mean) * 255.), 0, 255).astype(np.uint8)
      # 生成热图 默认是白色
      pred = debugger.gen_colormap(output['hm'][i].detach().cpu().numpy())
      # gt 热图
      gt = debugger.gen_colormap(batch['hm'][i].detach().cpu().numpy())
      #将预测热图和原图合并为一张图像，默认透明度比例为 3：7
      debugger.add_blend_img(img, pred, 'pred_hm')
      debugger.add_blend_img(img, gt, 'gt_hm')
      #将原图像放在图像字典中 作为增加预测框的预测图像
      debugger.add_img(img, img_id='out_pred')
      #遍历i为索引的检测结果
      for k in range(len(dets[i])):
        # 如果该检测score > 0.1
        if dets[i, k, 4] > opt.center_thresh:
          #添加对应的检测框 输出预测结果
          debugger.add_coco_bbox(dets[i, k, :4], dets[i, k, -1],
                                 dets[i, k, 4], img_id='out_pred')
      # 输出 正确的结果
      debugger.add_img(img, img_id='out_gt')
      for k in range(len(dets_gt[i])):
        # 在图像上绘制 gt 框，其实这里可以直接用 out_pred 对应的预测图，让其用不同的颜色即可，但是容易造成混淆
        if dets_gt[i, k, 4] > opt.center_thresh:
          debugger.add_coco_bbox(dets_gt[i, k, :4], dets_gt[i, k, -1],
                                 dets_gt[i, k, 4], img_id='out_gt')
      #将检测结果图像保存在Debug文件夹
      if opt.debug == 4:
        debugger.save_all_imgs(opt.debug_dir, prefix='{}'.format(iter_id))
      else:
      #否则会在屏幕显示图像，但是由于训练/推断速度快，所以图像就不会在屏幕出现
        debugger.show_all_imgs(pause=True)
   # 保存结果
   # output：模型输出 batch:批次 results:最终结果
  def save_result(self, output, batch, results):
    # 判断是否有偏移量回归值
    reg = output['reg'] if self.opt.reg_offset else None
    #对结果进行解码
    dets = ctdet_decode(
      output['hm'], output['wh'], reg=reg,
      cat_spec_wh=self.opt.cat_spec_wh, K=self.opt.K)
    #将其转化为cpu下的numpy
    dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])
    #进行后处理解码，并返回处理后的预测结果输出 [[x1,y1,x2,y2,score],...]
    dets_out = ctdet_post_process(
      dets.copy(), batch['meta']['c'].cpu().numpy(),
      batch['meta']['s'].cpu().numpy(),
      output['hm'].shape[2], output['hm'].shape[3], output['hm'].shape[1])
      # 将结果保存在结果字典中，键位图像id
    results[batch['meta']['img_id'].cpu().numpy()[0]] = dets_out[0]

3 加载数据

  print('Setting up data...')
  #加载验证集数据
  val_loader = torch.utils.data.DataLoader(
      Dataset(opt, 'val'), 
      batch_size=1, 
      shuffle=False,
      num_workers=1,
      pin_memory=True
  )
  #如果是 测试 则调用 val 方法
  if opt.test:
    _, preds = trainer.val(0, val_loader)
    val_loader.dataset.run_eval(preds, opt.save_dir)#进行结果评价
    return
  #加载训练集
  train_loader = torch.utils.data.DataLoader(
      Dataset(opt, 'train'), 
      batch_size=opt.batch_size, 
      shuffle=True,
      num_workers=opt.num_workers,
      pin_memory=True,
      drop_last=True
  )

在上面有个trainer.val方法，这个在下面进行说明

4 训练模型

  print('Starting training...')
  best = 1e10#设置损失函数哨兵 根据训练情况进行更新 
  #循环训练
  for epoch in range(start_epoch + 1, opt.num_epochs + 1):
    mark = epoch if opt.save_all else 'last'#是否每轮保存模型
    #训练
    log_dict_train, _ = trainer.train(epoch, train_loader)
    logger.write('epoch: {} |'.format(epoch))#写日志
    #获取训练的结果，不同的 loss 值
    for k, v in log_dict_train.items():
      logger.scalar_summary('train_{}'.format(k), v, epoch)
      logger.write('{} {:8f} | '.format(k, v))
    # 如果验证的间隔>0 或者 轮次对间隔取余为0 则保存模型
    if opt.val_intervals > 0 and epoch % opt.val_intervals == 0:
      save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(mark)), 
                 epoch, model, optimizer)
     #保存完模型，进行验证模型性能，让参数不具有梯度
      with torch.no_grad():
        # 返回预测的loss值 和 预测结果
        log_dict_val, preds = trainer.val(epoch, val_loader)
      #得到预测的数据
      for k, v in log_dict_val.items():
        logger.scalar_summary('val_{}'.format(k), v, epoch)
        logger.write('{} {:8f} | '.format(k, v))
      # 如果验证集的损失函数小于最好的情况，则将此模型保存为最好的
      if log_dict_val[opt.metric] < best:
        best = log_dict_val[opt.metric]
        save_model(os.path.join(opt.save_dir, 'model_best.pth'), 
                   epoch, model)
    else:
      #不然的话 保存最后一次模型
      save_model(os.path.join(opt.save_dir, 'model_last.pth'), 
                 epoch, model, optimizer)
    logger.write('\n')
    #如果需要更新学习率 则保存对应轮次 且更新学习率
    if epoch in opt.lr_step:
      save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(epoch)), 
                 epoch, model, optimizer)
      lr = opt.lr * (0.1 ** (opt.lr_step.index(epoch) + 1))
      print('Drop LR to', lr)
      #更新优化器的学习率
      for param_group in optimizer.param_groups:
          param_group['lr'] = lr
  logger.close()#关闭日志

以上函数里有 trainer.train() 和 trainer.val() 两个函数，这两个函数具体为：

  def val(self, epoch, data_loader):
    return self.run_epoch('val', epoch, data_loader)

  def train(self, epoch, data_loader):
    return self.run_epoch('train', epoch, data_loader)
  #关键的函数
  def run_epoch(self, phase, epoch, data_loader):
    # 得到带有损失函数的模型
    model_with_loss = self.model_with_loss
    if phase == 'train': #训练
      model_with_loss.train()
    else:#测试
      if len(self.opt.gpus) > 1:
        model_with_loss = self.model_with_loss.module
      model_with_loss.eval()
      torch.cuda.empty_cache()# 清空cuda的缓存

    opt = self.opt#获取配置信息
    results = {}#结果字典
    #记录和存储均值和现在的值
    data_time, batch_time = AverageMeter(), AverageMeter()
    #
    avg_loss_stats = {l: AverageMeter() for l in self.loss_stats}
    # 每轮的迭代次数
    num_iters = len(data_loader) if opt.num_iters < 0 else opt.num_iters
    #打印信息
    bar = Bar('{}/{}'.format(opt.task, opt.exp_id), max=num_iters)
    end = time.time()#计时
    #迭代数据加载器
    for iter_id, batch in enumerate(data_loader):
      if iter_id >= num_iters:
        break
      data_time.Eupdate(time.time() - end)#更新时间
      #迭代batchs数据 这个主要与所使用的数据加载器有关
      for k in batch:
        if k != 'meta':
          batch[k] = batch[k].to(device=opt.device, non_blocking=True)    
      # 模型的输出 loss 和 损失状态
      output, loss, loss_stats = model_with_loss(batch)
      loss = loss.mean()#得到损失均值
      if phase == 'train':#如果是训练需要反向传播更新参数
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
      batch_time.update(time.time() - end)#更新时间
      end = time.time()
      #控制台输出
      Bar.suffix = '{phase}: [{0}][{1}/{2}]|Tot: {total:} |ETA: {eta:} '.format(
        epoch, iter_id, num_iters, phase=phase,
        total=bar.elapsed_td, eta=bar.eta_td)
        
      for l in avg_loss_stats:
        avg_loss_stats[l].update(
          loss_stats[l].mean().item(), batch['input'].size(0))
        Bar.suffix = Bar.suffix + '|{} {:.4f} '.format(l, avg_loss_stats[l].avg)
      if not opt.hide_data_time:#在训练时，显示时间
        Bar.suffix = Bar.suffix + '|Data {dt.val:.3f}s({dt.avg:.3f}s) ' \
          '|Net {bt.avg:.3f}s'.format(dt=data_time, bt=batch_time)
      # 每隔几次打印一次，默认为每次打印
      if opt.print_iter > 0:
        if iter_id % opt.print_iter == 0:
          print('{}/{}| {}'.format(opt.task, opt.exp_id, Bar.suffix)) 
      else:
        bar.next()#每次打印
      # 如果debug>0 则处理结果 保存输出结果 一般在测试时候用
      if opt.debug > 0:
        self.debug(batch, output, iter_id)
      #默认为 False 保存所有的预测结果
      if opt.test:
        self.save_result(output, batch, results)
      del output, loss, loss_stats#进行垃圾回收
    
    bar.finish()
    #得到损失结果
    ret = {k: v.avg for k, v in avg_loss_stats.items()}
    #计算时间(分钟)
    ret['time'] = bar.elapsed_td.total_seconds() / 60.
    return ret, results

5 数据加载

这里以COCO数据集为例，\src\lib\datasets\dataset 下的 coco.py，如果是自己的数据集，可以将此复制，然后更改对应得部分。

class COCO(data.Dataset):
  num_classes = 80#数据集的种类数
  default_resolution = [512, 512]#送入网络的图像大小
  #数据集的均值和方差
  mean = np.array([0.40789654, 0.44719302, 0.47026115],
                   dtype=np.float32).reshape(1, 1, 3)
  std  = np.array([0.28863828, 0.27408164, 0.27809835],
                   dtype=np.float32).reshape(1, 1, 3)

  def __init__(self, opt, split):
    super(COCO, self).__init__()
    #数据的地址
    self.data_dir = os.path.join(opt.data_dir, 'coco')
    self.img_dir = os.path.join(self.data_dir, '{}2017'.format(split))
    #根据不同的情况加载数据的json文件
    if split == 'test':
      self.annot_path = os.path.join(
          self.data_dir, 'annotations', 
          'image_info_test-dev2017.json').format(split)
    else:
      if opt.task == 'exdet':
        self.annot_path = os.path.join(
          self.data_dir, 'annotations', 
          'instances_extreme_{}2017.json').format(split)
      else:
        self.annot_path = os.path.join(
          self.data_dir, 'annotations', 
          'instances_{}2017.json').format(split)
    self.max_objs = 128#最大对象数目
    #类名列表，注意第一个为 __background__ ,其余为正常的类名
    self.class_name = [
      '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
      'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
      'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
      'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
      'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
      'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
      'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass',
      'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
      'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
      'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv',
      'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
      'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
      'scissors', 'teddy bear', 'hair drier', 'toothbrush']
    #这个是类别索引，注意从 1 开始，不包括背景类，仅为有效类，这个是最容易搞混的，您的那个json文件的类别索引需要从 1 开始
    self._valid_ids = [
      1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 
      14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 
      24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 
      37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 
      48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 
      58, 59, 60, 61, 62, 63, 64, 65, 67, 70, 
      72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 
      82, 84, 85, 86, 87, 88, 89, 90]
     #生成种类字典
    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)#随机种子
    #特征变量 _eig_val 
    self._eig_val = np.array([0.2141788, 0.01817699, 0.00341571],
                             dtype=np.float32)
    #特征变量 _eig_vec 
    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 = split# 训练还是测试的指示变量
    self.opt = opt# 配置文件

    print('==> initializing coco 2017 {} data.'.format(split))
    self.coco = coco.COCO(self.annot_path)#加载数据集
    self.images = self.coco.getImgIds()#获取图像id
    self.num_samples = len(self.images)#得到图像数目

    print('Loaded {} {} samples'.format(split, self.num_samples))
  #转为小数点后两位
  def _to_float(self, 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:#遍历所有的图像id
      for cls_ind in all_bboxes[image_id]:#获得每个结果中类别索引
        #获得类别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为两位小数
          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))
          }
          #如果bbox的长度大于5 则说明结果中有极值点
          if len(bbox) > 5:
              extreme_points = list(map(self._to_float, bbox[5:13]))
              detection["extreme_points"] = extreme_points
          detections.append(detection)#追加检测结果
    return detections

  def __len__(self):
    return self.num_samples
  # 保存结果
  def 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()

6 预处理数据集

该文件在src\lib\datasets\sample文件夹下

class CTDetDataset(data.Dataset):
   # 将 (x1,y1,w,h)->(x1,y1,x2,y2)
  def _coco_box_to_bbox(self, box):
    bbox = np.array([box[0], box[1], box[0] + box[2], box[1] + box[3]],
                    dtype=np.float32)
    return bbox
  # 获取边界
  def _get_border(self, border, size):
    i = 1
    while size - border // i <= border // i:
        i *= 2
    return border // i
# 根据索引获取单独的一个图像数据
  def __getitem__(self, index):
    img_id = self.images[index]
    file_name = self.coco.loadImgs(ids=[img_id])[0]['file_name']
    img_path = os.path.join(self.img_dir, file_name)
    ann_ids = self.coco.getAnnIds(imgIds=[img_id])#加载标注文件
    anns = self.coco.loadAnns(ids=ann_ids)#得到标注信息
    num_objs = min(len(anns), self.max_objs)#最多可以加载128个对象
    img = cv2.imread(img_path)#读取图像
    # 图像高 宽
    height, width = img.shape[0], img.shape[1]
    # 图像中心坐标，
    c = np.array([img.shape[1] / 2., img.shape[0] / 2.], dtype=np.float32)
    if self.opt.keep_res:#如果保持原有的分辨率
      input_h = (height | self.opt.pad) + 1
      input_w = (width | self.opt.pad) + 1
      s = np.array([input_w, input_h], dtype=np.float32)
    else:# 获取 s 为img h 和 w的最大值
      s = max(img.shape[0], img.shape[1]) * 1.0
      #定义输入的高和宽为配置的高和宽(512)
      input_h, input_w = self.opt.input_h, self.opt.input_w
    # 是否翻折（False）
    flipped = False
    if self.split == 'train':# 如果是训练
       #进行随机翻转 not_rand_crop默认为 False
      if not self.opt.not_rand_crop:
        s = s * np.random.choice(np.arange(0.6, 1.4, 0.1))
        #获取宽边界
        w_border = self._get_border(128, img.shape[1])
        #获取高边界
        h_border = self._get_border(128, img.shape[0])
        # 取宽随机数
        c[0] = np.random.randint(low=w_border, high=img.shape[1] - w_border)
        # 取高随机数
        c[1] = np.random.randint(low=h_border, high=img.shape[0] - h_border)
      else:#否则是测试
        sf = self.opt.scale#放缩比例 默认 0.4
        cf = self.opt.shift# 应用平移扩增 默认 0.1
        #约束平移范围，得到平移后的中心坐标（在原图基础上平移 -02-0.2）
        c[0] += s * np.clip(np.random.randn()*cf, -2*cf, 2*cf)
        c[1] += s * np.clip(np.random.randn()*cf, -2*cf, 2*cf)
        #对图像大小进行缩放 (0.6-1.4)之间
        s = s * np.clip(np.random.randn()*sf + 1, 1 - sf, 1 + sf)
      #随机翻折
      if np.random.random() < self.opt.flip:
        flipped = True
        img = img[:, ::-1, :]
        c[0] =  width - c[0] - 1
    # 获取平移 缩放的变换矩阵
    trans_input = get_affine_transform(
      c, s, 0, [input_w, input_h])
    # 进行仿射变换 完成图像的平移缩放等
    inp = cv2.warpAffine(img, trans_input, 
                         (input_w, input_h),
                         flags=cv2.INTER_LINEAR)
    # 将图像像素值归一化
    inp = (inp.astype(np.float32) / 255.)
    # 如果是训练集 使用颜色扩增方法 对比度 亮度 饱和度
    if self.split == 'train' and not self.opt.no_color_aug:
      color_aug(self._data_rng, inp, self._eig_val, self._eig_vec)
    #对图像数据进行高斯归一化
    inp = (inp - self.mean) / self.std
    # 将图像转为 CHW 型
    inp = inp.transpose(2, 0, 1)
    # 输出的图像高 / 宽
    output_h = input_h // self.opt.down_ratio
    output_w = input_w // self.opt.down_ratio
    num_classes = self.num_classes#类别数目
    #输出图像的转换矩阵
    trans_output = get_affine_transform(c, s, 0, [output_w, output_h])
    # 热图 (94,128,128)
    hm = np.zeros((num_classes, output_h, output_w), dtype=np.float32)
    #(128，2)
    wh = np.zeros((self.max_objs, 2), dtype=np.float32)
    #(2,128,128)
    dense_wh = np.zeros((2, output_h, output_w), dtype=np.float32)
    reg = np.zeros((self.max_objs, 2), dtype=np.float32)#(128，2)
    ind = np.zeros((self.max_objs), dtype=np.int64)#(128,)
    reg_mask = np.zeros((self.max_objs), dtype=np.uint8)#(128,)
    cat_spec_wh = np.zeros((self.max_objs, num_classes * 2), dtype=np.float32)#(128,188)
    cat_spec_mask = np.zeros((self.max_objs, num_classes * 2), dtype=np.uint8)#(128,188)
    #使用 draw_umich_gaussian 确定高斯半径
    draw_gaussian = draw_msra_gaussian if self.opt.mse_loss else \
                    draw_umich_gaussian

    gt_det = []# gt信息
    for k in range(num_objs):
      ann = anns[k]# 索引为k的标注信息
      bbox = self._coco_box_to_bbox(ann['bbox'])#转换坐标
      cls_id = int(self.cat_ids[ann['category_id']])# 类别索引
      if flipped:#是否翻折 更改对应坐标值
        bbox[[0, 2]] = width - bbox[[2, 0]] - 1
      bbox[:2] = affine_transform(bbox[:2], trans_output)#转换坐标
      bbox[2:] = affine_transform(bbox[2:], trans_output)
      bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, output_w - 1)#约束坐标防止越界
      bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, output_h - 1)
      h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
      if h > 0 and w > 0:
        radius = gaussian_radius((math.ceil(h), math.ceil(w)))#生成高斯半径
        radius = max(0, int(radius))
        radius = self.opt.hm_gauss if self.opt.mse_loss else radius
        ct = np.array(
          [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32)#实例的中心点坐标
        ct_int = ct.astype(np.int32)
        draw_gaussian(hm[cls_id], ct_int, radius)#绘制高斯半径区域
        wh[k] = 1. * w, 1. * h#框宽和高
        ind[k] = ct_int[1] * output_w + ct_int[0]
        reg[k] = ct - ct_int#框偏移
        reg_mask[k] = 1#索引为k的类别mask为1
        cat_spec_wh[k, cls_id * 2: cls_id * 2 + 2] = wh[k]#索引为k对应种类的宽和高
        cat_spec_mask[k, cls_id * 2: cls_id * 2 + 2] = 1##索引为k对应种类的mask
        #默认为False，考虑全局的回归，为true时，在中心附近应用加权回归或仅在中心点应用回归
        if self.opt.dense_wh:
          draw_dense_reg(dense_wh, hm.max(axis=0), ct_int, wh[k], radius)
          #列表追加gt 信息[x1,y1,x2,y2,1,类别索引]
        gt_det.append([ct[0] - w / 2, ct[1] - h / 2, 
                       ct[0] + w / 2, ct[1] + h / 2, 1, cls_id])
    #gt 字典
    ret = {'input': inp, 'hm': hm, 'reg_mask': reg_mask, 'ind': ind, 'wh': wh}
    if self.opt.dense_wh:#默认False
      hm_a = hm.max(axis=0, keepdims=True)
      dense_wh_mask = np.concatenate([hm_a, hm_a], axis=0)
      ret.update({'dense_wh': dense_wh, 'dense_wh_mask': dense_wh_mask})
      del ret['wh']
    elif self.opt.cat_spec_wh:#默认False
      ret.update({'cat_spec_wh': cat_spec_wh, 'cat_spec_mask': cat_spec_mask})
      del ret['wh']
    if self.opt.reg_offset:#使用回归偏置
      ret.update({'reg': reg})
     #如果debug大于0且为训练集
    if self.opt.debug > 0 or not self.split == 'train':
      gt_det = np.array(gt_det, dtype=np.float32) if len(gt_det) > 0 else \
               np.zeros((1, 6), dtype=np.float32)
      meta = {'c': c, 's': s, 'gt_det': gt_det, 'img_id': img_id}
      ret['meta'] = meta
    return ret

总结

以上除了具体的网络结构没有详细的介绍，基本上把CenterNet代码的整个流程都串了一遍，如果有疑问的地方，欢迎讨论，如需转载，望注明出处！谢谢合作