COCO数据集格式、mask两种存储格式、官方包API详解

落日映苍穹つ 2024-05-08 06:40 5阅读 0赞

#### 文章目录 ####

*  前言
 *  COCO数据集格式
 *   *  总体组织结构
     *  \`image对象\`的主要属性：
     *  \`annotation对象\`的属性主要：
 *  pycocotools包中的COCO对象及API
 *   *  COCO对象的属性
     *  COCO对象的方法
 *  API源码分析
 *   *  mask保存的格式
     *  mask不同的保存格式如何进行可视化
     *  \`showAnns\`函数源码剖析

## 前言 ##

在前面的文章中我们介绍了如何使用COCO提供的官方API来进行coco格式数据集的可视化，下面我们主要来讲`COCO数据集格式`和`官方API实现`。

## COCO数据集格式 ##

coco数据集格式充分利用了面向对象的思路：整个标注文件是一个`json对象`，这个大的json对象包含几个主要的filed：`"info"`，`"licenses"`，`"categories"`，`"images"`，`"annotations"`。每个filed都是一个数组，里面包含所有的`image对象`和`annotation对象`。在coco格式中，每一张图片是一个json对象，每一个标注也是一个json对象，所有的对象都用一个`唯一的id`进行标识。注意，image对象和annotation对象的id是`分开来标识`的。

### 总体组织结构 ###

如图所示：  
![在这里插入图片描述][watermark_type_d3F5LXplbmhlaQ_shadow_50_text_Q1NETiBA6YGH5Yiw5Z2O5bCx5b6X6L-I6L-H5Y67_size_20_color_FFFFFF_t_70_g_se_x_16]

### `image对象`的主要属性： ###

*  id：image对象的id
 *  file\_name: 图片的文件名
 *  width: 图片的宽
 *  height: 图片的高

### `annotation对象`的属性主要： ###

*  id: 标注的id
 *  image\_id: 标注所属图片的id
 *  category\_id：此标注所属的类别
 *  iscrown：是否紧凑，默认为0
 *  area：标注的面积
 *  bbox：标注的框，左上角x坐标，左上角y坐标，框的宽，框的高
 *  segmentation：掩码的标记，有两种格式：多边形格式和RLE格式
 *  width: 标注所属图片的宽
 *  height: 标注所属图片的高  
    ![在这里插入图片描述][watermark_type_d3F5LXplbmhlaQ_shadow_50_text_Q1NETiBA6YGH5Yiw5Z2O5bCx5b6X6L-I6L-H5Y67_size_14_color_FFFFFF_t_70_g_se_x_16]

## pycocotools包中的COCO对象及API ##

为了方便可视化COCO格式的对象，官方为我们能提供了`pycocotools`这个工具包，下面我们来介绍一下这个包下的`COCO`类.

使用前通过`from pycocotools.coco import COCO`语句进行导入。

### COCO对象的属性 ###

*  self.dataset: 所加载的json文件
 *  self.anns：字典类型，保存所有的annotation对象，key是annotation的id，value是标注对象
 *  self.cats：字典类型，保存所有的catgory对象，key是catgory的id，value是类别对象
 *  self.imgs：字典类型，保存所有的image对象，key是image的id，value是图片对象
 *  self.imgToAnns：字典类型，保存所有的image到anno的对应关系，key是图片的id，value是数组对象，保存标注对象，`图片与标注是一对多的关系`
 *  self.catToImgs：字典类型，保存所有的cat到image的对应关系，key是类别的id，value是数组对象，保存图片对象，`类别与图片是一对多的关系`

### COCO对象的方法 ###

所有的方法如下图所示：  
![在这里插入图片描述][watermark_type_d3F5LXplbmhlaQ_shadow_50_text_Q1NETiBA6YGH5Yiw5Z2O5bCx5b6X6L-I6L-H5Y67_size_10_color_FFFFFF_t_70_g_se_x_16]

*  `getAnnoIds`：根据图片id或者类id来得到相应的标注id
 *  `getCatIds`: 得到类别id
 *  `getImgIds`: 根据类别id得到图片id
 *  `loadAnns`: 根据标注Id得到标注对象
 *  `loadImgs`: 根据图片ID得到图片对象
 *  `showAnnos`: 可视化掩码
 *  `loadRes`: 根据预测结果生成对应的COCO对象

总之，有了上面的这些方法，我们在`img`,`anno`, `catgory` 之间形成了一个闭环，三者可以互相转换，如下图：

![在这里插入图片描述][watermark_type_d3F5LXplbmhlaQ_shadow_50_text_Q1NETiBA6YGH5Yiw5Z2O5bCx5b6X6L-I6L-H5Y67_size_16_color_FFFFFF_t_70_g_se_x_16]

## API源码分析 ##

下面我们主要分析一下`showAnns`函数，看一下官方API是如何显示掩码的.

### mask保存的格式 ###

再分析源代码之前，我们首先想一下应该如何表示一个mask。mask就是一个不规则的分割结果，那么这个不规则的形状在计算机中应该如何保存呢？

主要有两种格式：`多边形表示法`，`RLE压缩表示法`

多边形表示法就是将`mask边缘的坐标`保存下来，类似于图一，COCO数据集中的mask就是以这种形式表示的；RLE是将mask所在的区域进行`二值化`后，压缩保存，类似于图二，预测结果中的mask就是以这种格式保存的。  
![在这里插入图片描述][1aaafb44587f4833a5cc906c17ec8ad8.png]  
图一  
![图二][5704db95eb80415bab3a53c8cb9db3f8.png]  
图二

### mask不同的保存格式如何进行可视化 ###

首先来说`多边形表示法`如何转化为二值掩膜。

对于多边形表示法，我们的思路是

*  1、先初始化一张和原图大小相等的mask，令所有的初值为零
 *  2、取出annos中多边形边缘点的坐标
 *  3、调用`cv2.fillPoly()`函数进行多边形填充

代码如下：

if 'segmentation' in ann:
       # 如果掩码为多边形格式
       if type(ann['segmentation']) == list:
           polys = []
           # 依次遍历掩码中多个多边形，一个掩膜中也可能包含多个多边形
           for seg in ann['segmentation']:
               poly = np.array(seg, dtype=np.int32).reshape((int(len(seg) / 2), 2))
               polys.append(poly)
           # 多边形填充, mask代表将要填充的初始图，polys存储所有的多边形边界坐标，1代表所填充的值
           cv2.fillPoly(mask, polys, 1)

再来说`RLE`格式如何转换为二值掩膜。

对于RLE格式的转化，我们可以直接借助官方的API`maskUtils.decode(rle)`函数实现，得到的对象直接就是我们想要的二值掩膜。

代码如下：

# 如果掩码为RLE格式
    else:
        if type(ann['segmentation']['counts']) == list:
            for seg in ann['segmentation']:
                poly = np.array(seg).reshape((int(len(seg) / 2), 2))
                # polygons.append(Polygon(poly))
                # color.append(c)
            # rle = maskUtils.frPyObjects([ann['segmentation']], t['height'], t['width'])
        else:
            rle = [ann['segmentation']]  # 针对预测结果而言，通常直接走这个分支
        m = maskUtils.decode(rle)
        # 解码后是相同的三通道，我们只取一个通道
        m = m[:, :, 0].reshape((m.shape[0], m.shape[1]))

完整的代码如下，这个函数可以根据image\_id和category\_id生成对应的二值掩膜：

def get_mask(coco: COCO, image_id, category_id):
    	"""
    	根据coco对象，image_id和cat_id生成掩膜
    	:param coco:
    	:param image_id:
    	:param category_id:
    	:return:
    	"""
    	image = coco.loadImgs(image_id)
    	image = image[0]  # 因为我们只会得到一张特定的图片
    	w, h = image["width"], image["height"]
    	# 每一类的掩膜
    	mask = np.zeros((h, w), dtype="uint8")
    	# 找到每一类的掩膜
    	anno_ids = coco.getAnnIds(imgIds=image_id, catIds=category_id, iscrowd=None)
    	if len(anno_ids) == 0:
    	    return mask
    	# 得到所有的标注，并依次处理每个标注
    	annos = coco.loadAnns(anno_ids)
    	for ann in annos:
    	    if 'segmentation' in ann:
    	        # 如果掩码为多边形格式
    	        if type(ann['segmentation']) == list:
    	            polys = []
    	            # 依次遍历掩码中多个多边形，一个掩膜中也可能包含多个多边形
    	            for seg in ann['segmentation']:
    	                poly = np.array(seg, dtype=np.int32).reshape((int(len(seg) / 2), 2))
    	                polys.append(poly)
    	            # 多边形填充
    	            cv2.fillPoly(mask, polys, 1)
    	
    	        # 如果掩码为RLE格式
    	        else:
    	            if type(ann['segmentation']['counts']) == list:
    	                for seg in ann['segmentation']:
    	                    poly = np.array(seg).reshape((int(len(seg) / 2), 2))
    	                    # polygons.append(Polygon(poly))
    	                    # color.append(c)
    	                # rle = maskUtils.frPyObjects([ann['segmentation']], t['height'], t['width'])
    	            else:
    	                rle = [ann['segmentation']]
    	            m = maskUtils.decode(rle)
    	            # 解码后是相同的三通道，我们只取一个通道
    	            m = m[:, :, 0].reshape((m.shape[0], m.shape[1]))
    	            idx = np.nonzero(m)
    	            mask[idx[0], idx[1]] = 1
    	return mask

### `showAnns`函数源码剖析 ###

好了，有了前面知识的铺垫，我们可以来看官方源码了。这里我把注释直接写进了代码中，同学们一定要耐心看完，这样才能有所收获~

def showAnns(self, anns, draw_bbox=False):
        """
        Display the specified annotations.
        :param anns (array of object): annotations to display
        :return: None
        """
        # 判断标注是否为空
        if len(anns) == 0:
            return 0
        # 判断当前任务是否为实例分割或目标检测
        if 'segmentation' in anns[0] or 'keypoints' in anns[0]:
            datasetType = 'instances'
        elif 'caption' in anns[0]:
            datasetType = 'captions'
        else:
            raise Exception('datasetType not supported')
        # 如果当前正在做实例分割或检测
        if datasetType == 'instances':
        	# 导入需要的包
            import matplotlib.pyplot as plt
            from matplotlib.collections import PatchCollection
            from matplotlib.patches import Polygon
    		# 获取当前活动的子图对象实例
            ax = plt.gca()
            ax.set_autoscale_on(False)
            polygons = []  # 定义数组，保存多边形掩码
            color = [] 
            for ann in anns:
                c = (np.random.random((1, 3))*0.6+0.4).tolist()[0]
                # 判断是否为分割任务
                if 'segmentation' in ann:
                	# 如果mask的保存格式为多边形表示法
                    if type(ann['segmentation']) == list:
                        # 遍历所有的segmentation
                        for seg in ann['segmentation']:
                        	# 获取多边形边界点的坐标
                            poly = np.array(seg).reshape((int(len(seg)/2), 2))
                            # 根据边界点坐标构造Polygon实例
                            polygons.append(Polygon(poly))
                            color.append(c)
                    # 如果mask的保存格式为RLE
                    else:
                        # 根据image_id得到image对象，主要是为了得到这张图片的宽和高
                        t = self.imgs[ann['image_id']]
                        # 如果是未经压缩处理的RLE格式
                        if type(ann['segmentation']['counts']) == list:
                            rle = maskUtils.frPyObjects([ann['segmentation']], t['height'], t['width'])
                        else:
                            rle = [ann['segmentation']]
                        # 解析rle得到二值掩膜
                        m = maskUtils.decode(rle)
                        img = np.ones( (m.shape[0], m.shape[1], 3) )
                        if ann['iscrowd'] == 1:
                            color_mask = np.array([2.0,166.0,101.0])/255
                        if ann['iscrowd'] == 0:
                            color_mask = np.random.random((1, 3)).tolist()[0]
                        for i in range(3):
                            img[:,:,i] = color_mask[i]
                        # 显示掩膜
                        ax.imshow(np.dstack( (img, m*0.5) ))
                if 'keypoints' in ann and type(ann['keypoints']) == list:
                    # turn skeleton into zero-based index
                    sks = np.array(self.loadCats(ann['category_id'])[0]['skeleton'])-1
                    kp = np.array(ann['keypoints'])
                    x = kp[0::3]
                    y = kp[1::3]
                    v = kp[2::3]
                    for sk in sks:
                        if np.all(v[sk]>0):
                            plt.plot(x[sk],y[sk], linewidth=3, color=c)
                    plt.plot(x[v>0], y[v>0],'o',markersize=8, markerfacecolor=c, markeredgecolor='k',markeredgewidth=2)
                    plt.plot(x[v>1], y[v>1],'o',markersize=8, markerfacecolor=c, markeredgecolor=c, markeredgewidth=2)
    
                if draw_bbox:
                    [bbox_x, bbox_y, bbox_w, bbox_h] = ann['bbox']
                    poly = [[bbox_x, bbox_y], [bbox_x, bbox_y+bbox_h], [bbox_x+bbox_w, bbox_y+bbox_h], [bbox_x+bbox_w, bbox_y]]
                    np_poly = np.array(poly).reshape((4,2))
                    polygons.append(Polygon(np_poly))
                    color.append(c)
    		# 将多边形放入多边形集合中，如果是RLE格式，那么多边形集合为空
            p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.4)
            # 叠加显示多边形集合
            ax.add_collection(p)
            # 为多边形画边界框
            p = PatchCollection(polygons, facecolor='none', edgecolors=color, linewidths=2)
            ax.add_collection(p)
        # 另一种任务，暂时用不到
        elif datasetType == 'captions':
            for ann in anns:
                print(ann['caption'])

[watermark_type_d3F5LXplbmhlaQ_shadow_50_text_Q1NETiBA6YGH5Yiw5Z2O5bCx5b6X6L-I6L-H5Y67_size_20_color_FFFFFF_t_70_g_se_x_16]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/08/d2dff7b279854ff0a8dedcbcb4a75895.png
[watermark_type_d3F5LXplbmhlaQ_shadow_50_text_Q1NETiBA6YGH5Yiw5Z2O5bCx5b6X6L-I6L-H5Y67_size_14_color_FFFFFF_t_70_g_se_x_16]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/08/cd1a9d493c5b4f28a4fc77275b9040ac.png
[watermark_type_d3F5LXplbmhlaQ_shadow_50_text_Q1NETiBA6YGH5Yiw5Z2O5bCx5b6X6L-I6L-H5Y67_size_10_color_FFFFFF_t_70_g_se_x_16]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/08/d453b39b84ea4ecd8a7b571db53f3109.png
[watermark_type_d3F5LXplbmhlaQ_shadow_50_text_Q1NETiBA6YGH5Yiw5Z2O5bCx5b6X6L-I6L-H5Y67_size_16_color_FFFFFF_t_70_g_se_x_16]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/08/5288e19c6f0149e4ba84c2664008082d.png
[1aaafb44587f4833a5cc906c17ec8ad8.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/08/d66ecc022ac747fd9a918e8196376eba.png
[5704db95eb80415bab3a53c8cb9db3f8.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/08/9f31ce55de4848bfa3a89f336f1a49c3.png