Supervisely 人像分割数据集格式转换

Supervisely 数据集的读取和解析有两种方法:1. 读取 json 文件自己解析;2. 使用 supervisely_lib 解析。

人像分割数据集中人像的标注有两种几何类型:位图(bitmap),多边形(polygon)。自己解析需要对这两种类型进行处理,参考 “Supervisely 人像分割数据集格式解析”。个人建议使用方法2。supervisely_lib 提供了统一的API,你不需要关心几何类型,它们(位图,多边形)都可以使用相同的API获取 mask,轮廓,box。

这里以生成 Mask 图像为例。

一、安装 supervisely_lib

supervisely_lib 项目地址:https://github.com/supervisely/supervisely

安装方法参见项目主页(README.md)。

安装过程中会依赖一些第三方库,其中要求 opencv-python 版本小于4.0.0。因为我的Python环境中已安装 opencv-python 4.2.0,所以我对 supervisely 稍微的修改,差异如下:

diff --git a/setup.py b/setup.py
index dad24a3..4cc0698 100644
--- a/setup.py
+++ b/setup.py
@@ -25,7 +25,7 @@ setup(
         "jsonschema>=2.6.0,<3.0.0",
         "matplotlib>=3.0.0",
         "numpy>=1.14.3",
-        "opencv-python>=3.4.1,<4.0.0",
+        "opencv-python>=3.4.1",
         "pandas>=1.0.3",
         "pascal-voc-writer>=0.1.4",
         "PTable>=0.9.2",
diff --git a/supervisely_lib/geometry/bitmap.py b/supervisely_lib/geometry/bitmap.py
index 30880b8..ab51919 100644
--- a/supervisely_lib/geometry/bitmap.py
+++ b/supervisely_lib/geometry/bitmap.py
@@ -117,7 +117,7 @@ class Bitmap(BitmapBase):
         self.to_bbox().get_cropped_numpy_slice(bitmap)[self.data] = color
 
     def _draw_contour_impl(self, bitmap, color, thickness=1, config=None):
-        _, contours, _ = cv2.findContours(self.data.astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+        contours, _ = cv2.findContours(self.data.astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
         if contours is not None:
             for cont in contours:
                 cont[:, :] += (self.origin.col, self.origin.row)  # cont with shape (rows, ?, 2)

二、准备数据集

数据集下载后解压并重命名,这里重命名为person。

数据集的下载,参考:https://blog.csdn.net/SimleCat/article/details/107021152

我的目录结构如下:

├── code
├── person
│   ├── ds#
│   │   ├── ann
│   │   └── img
│  └── meta.json
└── person_mask

:"#" 取[1, 13],code 和 person_mask 是我自己创建的。

注意:person 为数据集目录,该目录下不可创建新的目录,否则 supervisely_lib 加载 Project 时会报错。

三、生成 mask

code 目录下新建文件 gen_mask.py,文件内容如下:

import os
import numpy as np
import cv2
import tqdm
import supervisely_lib as sly
from matplotlib import pyplot as plt


def genMask(dataset_root_dir="./person",
            output_dir="./person_mask",
            label_color=[255,255,255]):
    ''' 生成Mask
        dataset_root_dir: 数据集根目录
        output_dir: 生成mask的保存目录
        label_color: mask 标注的颜色[R, G, B]
    '''
    project = sly.Project(dataset_root_dir, sly.OpenMode.READ)

    # 打印数据集信息
    print(f"Project name     : {project.name}")
    print(f"Project directory: {project.directory}")
    print(f"Total images     : {project.total_items}")
    print(f"Dataset names    : {project.datasets.keys()}")
    print()
    print(project.meta)

    pbar = tqdm.tqdm(total=project.total_items)
    for dataset in project:
        for item_name in dataset:
            # 更新进度条
            pbar.update(1)
            # 获取原始图像和标注文件路径
            item_paths = dataset.get_item_paths(item_name)
            # 加载注释文件
            ann = sly.Annotation.load_json_file(item_paths.ann_path, project.meta)
            # 创建一个用于渲染标注的3通道黑色画布,
            ann_render = np.zeros(ann.img_size + (3,), dtype=np.uint8)
            # 渲染所有的标注(该数据集只有人)
            ann.draw(ann_render, color=label_color)
            # ann_render shape: (h, w, c), pixel: (R, G, B)
            # RGB -> BGR,用于Opencv
            ann_render = ann_render[..., ::-1]
            # mask 保存目录不存在,则创建
            save_dir = os.path.join(output_dir, dataset.name)
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            # 保存mask
            mask_path = os.path.join(save_dir, item_name)
            cv2.imwrite(mask_path, ann_render)
    pbar.close()


if __name__ == '__main__':
    import fire
    fire.Fire(genMask, name='gen_mask')
    '''
    Usage:
    # 1. 使用默认配置
    python code/gen_mask.py
    # 2. 指定数据库目录和mask输出目录
    python gen_mask.py --dataset_root_dir ../person --output_dir ../person_mask
    '''

gen_mask.py 使用方法,请参考上面代码的最后 "Usage" 部分。

supervisely_lib 更多的使用示例,参考:https://github.com/supervisely/supervisely/blob/master/help/jupyterlab_scripts/src/tutorials/01_project_structure/project.ipynb

你可能感兴趣的:(数据集,深度学习)