CenterNet学习记录(一)——COCO数据处理

1、前言

本个CenterNet系列的代码地址为:zzzxxxttt/torch_simple_CenterNet_45
此次主要记录的是数据处理方面,主要针对COCO数据。主要涉及脚本为:
datasets/coco.pyutils/images.py
以下主要结合代码进行对一些关键点和难点进行记录。

2、主要内容

CenterNet学习记录(一)——COCO数据处理_第1张图片
这里主要对COCO数据集包含的对象进行罗列,这里额外加入了__background__,代表背景,同时列出对象所对应的编号,但其并不是连续的,从末尾是90而类别是80+1即可知道不对应,所以其并不连续;然后就是一些数据处理所需要用到的一些特定的常量。(图在上方 > - <)

def __init__(self, data_dir, split, split_ratio=1.0, img_size=512):    	
	super(COCO, self).__init__()    
	self.num_classes = 80    
	self.class_name = COCO_NAMES    
	self.valid_ids = COCO_IDS
	   
	self.cat_ids = {v: i for i, v in enumerate(self.valid_ids)}
	
	self.data_rng = np.random.RandomState(123)    
	self.eig_val = np.array(COCO_EIGEN_VALUES, dtype=np.float32)    
	self.eig_vec = np.array(COCO_EIGEN_VECTORS, dtype=np.float32)    
	self.mean = np.array(COCO_MEAN, dtype=np.float32)[None, None, :]    
	self.std = np.array(COCO_STD, dtype=np.float32)[None, None, :]
	
	   	
	
	 	self.split = split   # 'train'  'val'    
  	self.data_dir = os.path.join(data_dir, 'COCO2017/')    
  	self.img_dir = os.path.join(self.data_dir, 'images/%s2017' % split)
          
   	if split == 'test':      
    		self.annot_path = os.path.join(self.data_dir, 'annotations', 'image_info_test-dev2017.json')    
   	else:      
   		self.annot_path = os.path.join(self.data_dir, 'annotations', 'instances_%s2017.json' % split)
   		
    	self.max_objs = 128    
    	self.padding = 127  
    	self.down_ratio = 4
    	self.img_size = {'h': img_size, 'w': img_size}
    	self.fmap_size = {'h': img_size // self.down_ratio, 'w': img_size // self.down_ratio}
    	self.rand_scales = np.arange(0.6, 1.4, 0.1)
    	self.gaussian_iou = 0.7
    
    	print('==> initializing coco 2017 %s data.' % split)    
    	self.coco = coco.COCO(self.annot_path)
    	self.images = self.coco.getImgIds()

接着是class COCO(data.Dataset)这个类,这是这个脚本的核心。def __init__()是进行的一些初始化处理操作,其中比较重要的是:
self.annot_path = os.path.join(self.data_dir, 'annotations', 'instances_%s2017.json' % split)
self.coco = coco.COCO(self.annot_path)
self.images = self.coco.getImgIds()

首先我们看coco这个库,需要我们导入pycocotools,这个的安装,具体可参考如下博客:
Windows下安装 pycocotools
安装好这个后,我们便来获取coco数据,如下是coco2017的目录结构:
CenterNet学习记录(一)——COCO数据处理_第2张图片
训练所采用的是instances_train2017.json,通过此json标签文件,加上coco库,便可轻松获取COCO数据集我们想要的内容,如图片,图片所包含的物体的种类和编号,已经对应的bbox,都可以,具体操作流程如下:

self.coco = coco.COCO(self.annot_path)
self.images = self.coco.getImgIds()
#这里以index=0为例
img_id = self.images[0]
img_path = os.path.join(self.img_dir, self.coco.loadImgs(ids=[img_id])[0]['file_name'])
ann_ids = self.coco.getAnnIds(imgIds=[img_id])
annotations = self.coco.loadAnns(ids=ann_ids)

labels = np.array([self.cat_ids[anno['category_id']] for anno in annotations])
bboxes = np.array([anno['bbox'] for anno in annotations], dtype=np.float32)

下图是获取到的self.coco的内容:
CenterNet学习记录(一)——COCO数据处理_第3张图片
这里的anns,是图片中物体所在位置的bbox的坐标,其下的category_id,代表物体所属种类的编号,id,代表所属图片的标号,由于一张图片中可能存在多个物体,故可能出现不同的ann所属id相同的情况。这里的长度是860001。
CenterNet学习记录(一)——COCO数据处理_第4张图片
这个的imgs,主要的作用是提供img所在路径,即提供file_name,图片名称,通过图片名称和根目录来读取图片。其长度为118287。
CenterNet学习记录(一)——COCO数据处理_第5张图片
然后是def __getitem__(self, index)这个函数,也是数据处理和获取的主要函数。
下面这些主要是获取每张图片上的标签信息,包括其所属类别,bbox框,同时最后转化为[x1,y1,x2,y2]这种格式。(需要记住这种获取coco数据集的方法)

  def __getitem__(self, index):

    img_id = self.images[index]
    img_path = os.path.join(self.img_dir, self.coco.loadImgs(ids=[img_id])[0]['file_name'])
    #print(img_path)
    ann_ids = self.coco.getAnnIds(imgIds=[img_id])
    annotations = self.coco.loadAnns(ids=ann_ids)

    labels = np.array([self.cat_ids[anno['category_id']] for anno in annotations])
    bboxes = np.array([anno['bbox'] for anno in annotations], dtype=np.float32)

    if len(bboxes) == 0:
      bboxes = np.array([[0., 0., 0., 0.]], dtype=np.float32)
      labels = np.array([[0]])
    bboxes[:, 2:] += bboxes[:, :2]  # xywh to xyxy

然后是数据增强处理(代码见下方,示意图下),包括翻转、仿射变化(包括放缩,裁剪等)。翻转,即flipped为True时,代表进行翻转,img = img[ :, : :-1 , : ]作用是水平翻转,img = img[ : : -1 , : , : ]代表的是上下翻转;然后是get_border()这个函数,位于utils/image.py中,作用是当图片的w/h大于256时,便返回128,否则返回64,意思是图片较大时返回128,较小是返回64,然后center便是其选定的中心点,可见下面的示意图。然后是get_affine_transform()这个函数,其作用是求warpAffine()所需的变换矩阵,求变换矩阵,至少需要3个点,刚才的center是一个点,然后通过旋转一定角度又是一个点,在通过对称又可以获得一个点。这些点的处理全部位于utils/image.py中的get_affine_transform()函数中,其中,get_dir()是获取旋转后的点,get_3rd_point是获取对称的点。
CenterNet学习记录(一)——COCO数据处理_第6张图片

flipped = False
    if self.split == 'train':
      scale = scale * np.random.choice(self.rand_scales)
      w_border = get_border(128, width)
      h_border = get_border(128, height)

      center[0] = np.random.randint(low=w_border, high=width - w_border)
      center[1] = np.random.randint(low=h_border, high=height - h_border)

      if np.random.random() < 0.5:
        flipped = True
        img = img[:, ::-1, :]
        center[0] = width - center[0] - 1

    trans_img = get_affine_transform(center, scale, 0, [self.img_size['w'], self.img_size['h']])
    img = cv2.warpAffine(img, trans_img, (self.img_size['w'], self.img_size['h']))

接着,是下面的代码,归一化操作,颜色增强,均值处理,通道变化,特征图的变化矩阵,初始化所要用的参数,对bbox坐标点也进行相应的仿射变化,然后求出包含对象的中心点。

	img = img.astype(np.float32) / 255.

    if self.split == 'train':
      color_aug(self.data_rng, img, self.eig_val, self.eig_vec)

    img -= self.mean
    img /= self.std
    img = img.transpose(2, 0, 1)  # from [H, W, C] to [C, H, W]

    trans_fmap = get_affine_transform(center, scale, 0, [self.fmap_size['w'], self.fmap_size['h']])

    hmap = np.zeros((self.num_classes, self.fmap_size['h'], self.fmap_size['w']), dtype=np.float32)  # heatmap
    w_h_ = np.zeros((self.max_objs, 2), dtype=np.float32)  # width and height
    regs = np.zeros((self.max_objs, 2), dtype=np.float32)  # regression
    inds = np.zeros((self.max_objs,), dtype=np.int64)
    ind_masks = np.zeros((self.max_objs,), dtype=np.uint8)

    # detections = []
    for k, (bbox, label) in enumerate(zip(bboxes, labels)):
      if flipped:
        bbox[[0, 2]] = width - bbox[[2, 0]] - 1
      bbox[:2] = affine_transform(bbox[:2], trans_fmap)
      bbox[2:] = affine_transform(bbox[2:], trans_fmap)
      bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.fmap_size['w'] - 1)
      bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.fmap_size['h'] - 1)
      h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
      if h > 0 and w > 0:
        obj_c = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32)
        obj_c_int = obj_c.astype(np.int32)

最后,是关于高斯处理的一些知识,具体见下列代码。首先是获取’‘高斯半径’’,具体求解见gaussian_radius这个函数,见下列图片,可以看出,是r1,r2,r3中的最小值。当时初看的时候,感觉很像二次函数的求根公式,后来经查询资料,果然真是。。。具体推导,可见下列我的手写版,也可参考下列网址讲解heatmap里面如何应用高斯散射核。求出高斯半径之后,便是绘制高斯核,这个主要是要了解高斯的二位表达式,然后数据获取到此基本结束,最后return所需要的东西即可。
CenterNet学习记录(一)——COCO数据处理_第7张图片

		radius = max(0, int(gaussian_radius((math.ceil(h), math.ceil(w)), self.gaussian_iou)))
        draw_umich_gaussian(hmap[label], obj_c_int, radius)
        w_h_[k] = 1. * w, 1. * h
        regs[k] = obj_c - obj_c_int  # discretization error
        inds[k] = obj_c_int[1] * self.fmap_size['w'] + obj_c_int[0]
        ind_masks[k] = 1

CenterNet学习记录(一)——COCO数据处理_第8张图片

3、总结

这部分代码不是很难,主要是对一些数据处理方法、数学知识等的理解,明白代码所代表的的含义,理解原理才是难点。故记录一下。

你可能感兴趣的:(cv)