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

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

这里主要对COCO数据集包含的对象进行罗列&＃xff0c;这里额外加入了__background__&＃xff0c;代表背景&＃xff0c;同时列出对象所对应的编号&＃xff0c;但其并不是连续的&＃xff0c;从末尾是90而类别是80&＃43;1即可知道不对应&＃xff0c;所以其并不连续&＃xff1b;然后就是一些数据处理所需要用到的一些特定的常量。&＃xff08;图在上方 > - <&＃xff09;

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

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

首先我们看coco这个库&＃xff0c;需要我们导入pycocotools&＃xff0c;这个的安装&＃xff0c;具体可参考如下博客&＃xff1a;
Windows下安装 pycocotools
安装好这个后&＃xff0c;我们便来获取coco数据&＃xff0c;如下是coco2017的目录结构&＃xff1a;

训练所采用的是instances_train2017.json&＃xff0c;通过此json标签文件&＃xff0c;加上coco库&＃xff0c;便可轻松获取COCO数据集我们想要的内容&＃xff0c;如图片&＃xff0c;图片所包含的物体的种类和编号&＃xff0c;已经对应的bbox&＃xff0c;都可以&＃xff0c;具体操作流程如下&＃xff1a;

self.coco &＃61; coco.COCO(self.annot_path)
self.images &＃61; self.coco.getImgIds()
#这里以index&＃61;0为例
img_id &＃61; self.images[0]
img_path &＃61; os.path.join(self.img_dir, self.coco.loadImgs(ids&＃61;[img_id])[0][&＃39;file_name&＃39;])
ann_ids &＃61; self.coco.getAnnIds(imgIds&＃61;[img_id])
annotations &＃61; self.coco.loadAnns(ids&＃61;ann_ids)labels &＃61; np.array([self.cat_ids[anno[&＃39;category_id&＃39;]] for anno in annotations])
bboxes &＃61; np.array([anno[&＃39;bbox&＃39;] for anno in annotations], dtype&＃61;np.float32)


下图是获取到的self.coco的内容&＃xff1a;

这里的anns&＃xff0c;是图片中物体所在位置的bbox的坐标&＃xff0c;其下的category_id&＃xff0c;代表物体所属种类的编号&＃xff0c;id&＃xff0c;代表所属图片的标号&＃xff0c;由于一张图片中可能存在多个物体&＃xff0c;故可能出现不同的ann所属id相同的情况。这里的长度是860001。

这个的imgs&＃xff0c;主要的作用是提供img所在路径&＃xff0c;即提供file_name&＃xff0c;图片名称&＃xff0c;通过图片名称和根目录来读取图片。其长度为118287。

然后是def __getitem__(self, index)这个函数&＃xff0c;也是数据处理和获取的主要函数。
下面这些主要是获取每张图片上的标签信息&＃xff0c;包括其所属类别&＃xff0c;bbox框&＃xff0c;同时最后转化为[x1,y1,x2,y2]这种格式。(需要记住这种获取coco数据集的方法)

def __getitem__(self, index):img_id &＃61; self.images[index]img_path &＃61; os.path.join(self.img_dir, self.coco.loadImgs(ids&＃61;[img_id])[0][&＃39;file_name&＃39;])#print(img_path)ann_ids &＃61; self.coco.getAnnIds(imgIds&＃61;[img_id])annotations &＃61; self.coco.loadAnns(ids&＃61;ann_ids)labels &＃61; np.array([self.cat_ids[anno[&＃39;category_id&＃39;]] for anno in annotations])bboxes &＃61; np.array([anno[&＃39;bbox&＃39;] for anno in annotations], dtype&＃61;np.float32)if len(bboxes) &＃61;&＃61; 0:bboxes &＃61; np.array([[0., 0., 0., 0.]], dtype&＃61;np.float32)labels &＃61; np.array([[0]])bboxes[:, 2:] &＃43;&＃61; bboxes[:, :2] # xywh to xyxy


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

flipped &＃61; Falseif self.split &＃61;&＃61; &＃39;train&＃39;:scale &＃61; scale * np.random.choice(self.rand_scales)w_border &＃61; get_border(128, width)h_border &＃61; get_border(128, height)center[0] &＃61; np.random.randint(low&＃61;w_border, high&＃61;width - w_border)center[1] &＃61; np.random.randint(low&＃61;h_border, high&＃61;height - h_border)if np.random.random() < 0.5:flipped &＃61; Trueimg &＃61; img[:, ::-1, :]center[0] &＃61; width - center[0] - 1trans_img &＃61; get_affine_transform(center, scale, 0, [self.img_size[&＃39;w&＃39;], self.img_size[&＃39;h&＃39;]])img &＃61; cv2.warpAffine(img, trans_img, (self.img_size[&＃39;w&＃39;], self.img_size[&＃39;h&＃39;]))


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

img &＃61; img.astype(np.float32) / 255.if self.split &＃61;&＃61; &＃39;train&＃39;:color_aug(self.data_rng, img, self.eig_val, self.eig_vec)img -&＃61; self.meanimg /&＃61; self.stdimg &＃61; img.transpose(2, 0, 1) # from [H, W, C] to [C, H, W]trans_fmap &＃61; get_affine_transform(center, scale, 0, [self.fmap_size[&＃39;w&＃39;], self.fmap_size[&＃39;h&＃39;]])hmap &＃61; np.zeros((self.num_classes, self.fmap_size[&＃39;h&＃39;], self.fmap_size[&＃39;w&＃39;]), dtype&＃61;np.float32) # heatmapw_h_ &＃61; np.zeros((self.max_objs, 2), dtype&＃61;np.float32) # width and heightregs &＃61; np.zeros((self.max_objs, 2), dtype&＃61;np.float32) # regressioninds &＃61; np.zeros((self.max_objs,), dtype&＃61;np.int64)ind_masks &＃61; np.zeros((self.max_objs,), dtype&＃61;np.uint8)# detections &＃61; []for k, (bbox, label) in enumerate(zip(bboxes, labels)):if flipped:bbox[[0, 2]] &＃61; width - bbox[[2, 0]] - 1bbox[:2] &＃61; affine_transform(bbox[:2], trans_fmap)bbox[2:] &＃61; affine_transform(bbox[2:], trans_fmap)bbox[[0, 2]] &＃61; np.clip(bbox[[0, 2]], 0, self.fmap_size[&＃39;w&＃39;] - 1)bbox[[1, 3]] &＃61; np.clip(bbox[[1, 3]], 0, self.fmap_size[&＃39;h&＃39;] - 1)h, w &＃61; bbox[3] - bbox[1], bbox[2] - bbox[0]if h > 0 and w > 0:obj_c &＃61; np.array([(bbox[0] &＃43; bbox[2]) / 2, (bbox[1] &＃43; bbox[3]) / 2], dtype&＃61;np.float32)obj_c_int &＃61; obj_c.astype(np.int32)


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

radius &＃61; 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] &＃61; 1. * w, 1. * hregs[k] &＃61; obj_c - obj_c_int # discretization errorinds[k] &＃61; obj_c_int[1] * self.fmap_size[&＃39;w&＃39;] &＃43; obj_c_int[0]ind_masks[k] &＃61; 1


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