【深度学习】CenterNet-better-plus代码解读

CenterNet-better-plus代码解读

源代码 https://github.com/lbin/CenterNet-better-plus

detectron2架构

gt生成部分

centernet/centernet_gt.py

首先看一下经过处理的inputs

instance类里面包含了图片的gt_bbox、gt_classes等信息

import numpy as np
import torch


class CenterNetGT(object):
    @staticmethod
    def generate(config, batched_input):
        
        box_scale = 1 / config.MODEL.CENTERNET.DOWN_SCALE
        num_classes = config.MODEL.CENTERNET.NUM_CLASSES
        output_size = config.MODEL.CENTERNET.OUTPUT_SIZE #[128,128]
        min_overlap = config.MODEL.CENTERNET.MIN_OVERLAP # 0.7
        tensor_dim = config.MODEL.CENTERNET.TENSOR_DIM # 128

        scoremap_list, wh_list, reg_list, reg_mask_list, index_list = [[] for i in range(5)] #初始化gt列表
        
        for data in batched_input:
            # img_size = (data['height'], data['width'])

            bbox_dict = data["instances"].get_fields() # 字典形式保存了gt_bbox和gt_classes

            # init gt tensors
            gt_scoremap = torch.zeros(num_classes, *output_size) # torch.Size([80, 128, 128]) cpu
            gt_wh = torch.zeros(tensor_dim, 2) # torch.Size([128, 2])
            gt_reg = torch.zeros_like(gt_wh) # torch.Size([128, 2]) 中心点的偏移量
            reg_mask = torch.zeros(tensor_dim) # torch.Size([128])
            gt_index = torch.zeros(tensor_dim) # torch.Size([128])
            # pass

            boxes, classes = bbox_dict["gt_boxes"], bbox_dict["gt_classes"]
            num_boxes = boxes.tensor.shape[0]
            boxes.scale(box_scale, box_scale)