将CityScapes数据集转换为COCO格式的实例分割数据集

##¥##2019年结束前最后两小时,写出本年度最后一篇博客,不是博主没地方去跨年,实在是Coding乐趣无穷ahhhh ##¥##

COCO数据集是研究实例分割的通用数据集,但是针对道路环境的数据较少,对于自动驾驶和移动机器人的应用并不适配。而提供道路环境的实例分割数据的数据集目前有KITTI、CityScapes、Apllo等,CityScapes算是应用较广,数据质量较高的一个,所以,本篇博文给大家分享一下将CityScapes的实例标注转换为COCO格式的过程。

首先,COCO的标注信息格式参考博文:https://blog.csdn.net/u012609509/article/details/88680841,COCO的图片文件均放在一个文件夹下,而其标注信息都放在annotation文件夹的不同.json文本中,一个文本中包含了对所有图片以及某一任务的所有图片的标注信息。

CityScapes提供的标注信息有彩色语义图、实例id单通道图、语义id单通道图以及所有语义polygon的文本描述.json:

一个COCO标注文件中包含info、lisenses、categories、images以及annotations,代码初始预先设定了前三项,categories可以根据自己的需要进行修改,我保留的类别有5类,car、pedestrian(person)、truck、bus、rider,对应cityscapes中类别id的24-28, 以下是cityscapes完整的类别定义:

List of cityscapes labels:

                     name |  id | trainId |       category | categoryId | hasInstances | ignoreInEval
    --------------------------------------------------------------------------------------------------
                unlabeled |   0 |     255 |           void |          0 |            0 |            1
              ego vehicle |   1 |     255 |           void |          0 |            0 |            1
     rectification border |   2 |     255 |           void |          0 |            0 |            1
               out of roi |   3 |     255 |           void |          0 |            0 |            1
                   static |   4 |     255 |           void |          0 |            0 |            1
                  dynamic |   5 |     255 |           void |          0 |            0 |            1
                   ground |   6 |     255 |           void |          0 |            0 |            1
                     road |   7 |       0 |           flat |          1 |            0 |            0
                 sidewalk |   8 |       1 |           flat |          1 |            0 |            0
                  parking |   9 |     255 |           flat |          1 |            0 |            1
               rail track |  10 |     255 |           flat |          1 |            0 |            1
                 building |  11 |       2 |   construction |          2 |            0 |            0
                     wall |  12 |       3 |   construction |          2 |            0 |            0
                    fence |  13 |       4 |   construction |          2 |            0 |            0
               guard rail |  14 |     255 |   construction |          2 |            0 |            1
                   bridge |  15 |     255 |   construction |          2 |            0 |            1
                   tunnel |  16 |     255 |   construction |          2 |            0 |            1
                     pole |  17 |       5 |         object |          3 |            0 |            0
                polegroup |  18 |     255 |         object |          3 |            0 |            1
            traffic light |  19 |       6 |         object |          3 |            0 |            0
             traffic sign |  20 |       7 |         object |          3 |            0 |            0
               vegetation |  21 |       8 |         nature |          4 |            0 |            0
                  terrain |  22 |       9 |         nature |          4 |            0 |            0
                      sky |  23 |      10 |            sky |          5 |            0 |            0
                   person |  24 |      11 |          human |          6 |            1 |            0
                    rider |  25 |      12 |          human |          6 |            1 |            0
                      car |  26 |      13 |        vehicle |          7 |            1 |            0
                    truck |  27 |      14 |        vehicle |          7 |            1 |            0
                      bus |  28 |      15 |        vehicle |          7 |            1 |            0
                  caravan |  29 |     255 |        vehicle |          7 |            1 |            1
                  trailer |  30 |     255 |        vehicle |          7 |            1 |            1
                    train |  31 |      16 |        vehicle |          7 |            1 |            0
               motorcycle |  32 |      17 |        vehicle |          7 |            1 |            0
                  bicycle |  33 |      18 |        vehicle |          7 |            1 |            0
            license plate |  -1 |      -1 |        vehicle |          7 |            0 |            1

具体的转换代码如下:

import cv2
import numpy as np
import os, glob
import datetime
import json
import os
import re
import fnmatch
from PIL import Image
import numpy as np
from pycococreatortools import pycococreatortools
 
ROOT_DIR = '/data/cityscapes/val'
IMAGE_DIR = os.path.join(ROOT_DIR, "images/frankfurt")
ANNOTATION_DIR = os.path.join(ROOT_DIR, "gt/frankfurt")
INSTANCE_DIR = os.path.join(ROOT_DIR, "instances") 

INFO = {
    "description": "Cityscapes_Instance Dataset",
    "url": "https://github.com/waspinator/pycococreator",
    "version": "0.1.0",
    "year": "2019",
    "contributor": "Kevin_Jia",
    "date_created": "2019-12-30 16:16:16.123456"
}
 
LICENSES = [
    {
        "id": 1,
        "name": "Attribution-NonCommercial-ShareAlike License",
        "url": "http://creativecommons.org/licenses/by-nc-sa/2.0/"
    }
]

CATEGORIES = [
    {
        'id': 1,
        'name': 'car',
        'supercategory': 'cityscapes',
    },
        {
        'id': 2,
        'name': 'pedestrian',
        'supercategory': 'cityscapes',
    },
    {
        'id': 3,
        'name': 'truck',
        'supercategory': 'cityscapes',
    },
    {
        'id': 4,
        'name': 'bus',
        'supercategory': 'cityscapes',
    },
    {
        'id': 5,
        'name': 'rider',
        'supercategory': 'cityscapes',
    }
]

background_label = list(range(-1, 24, 1)) + list(range(29, 34, 1))

def masks_generator(imges):
    idx = 0
    for pic_name in imges:
        annotation_name = pic_name.split('_')[0] + '_' + pic_name.split('_')[1] + '_' + pic_name.split('_')[2] + '_gtFine_instanceIds.png'
        print(annotation_name)
        annotation = cv2.imread(os.path.join(ANNOTATION_DIR, annotation_name), -1)
        name = pic_name.split('.')[0]
        h, w = annotation.shape[:2]
        ids = np.unique(annotation)
        for id in ids:
            if id in background_label:
                continue
            instance_id = id
            class_id = instance_id // 1000
            if class_id == 24:
                instance_class = 'pedestrian'
            elif class_id == 25:
                instance_class = 'rider' 
            elif class_id == 26:
                instance_class = 'car'
            elif class_id == 27:
                instance_class = 'truck'
            elif class_id == 28:
                instance_class = 'bus'
            else:
                continue
            print(instance_id)
            instance_mask = np.zeros((h, w, 3),dtype=np.uint8)
            mask = annotation == instance_id
            instance_mask[mask] = 255
            mask_name = name + '_' + instance_class + '_' + str(idx) + '.png'
            cv2.imwrite(os.path.join(INSTANCE_DIR, mask_name), instance_mask)
            idx += 1

def filter_for_pic(files):
    file_types = ['*.jpeg', '*.jpg', '*.png']
    file_types = r'|'.join([fnmatch.translate(x) for x in file_types])
    files = [f for f in files if re.match(file_types, f)]
    # files = [os.path.join(root, f) for f in files]
    return files
 
def filter_for_instances(root, files, image_filename):
    file_types = ['*.png']
    file_types = r'|'.join([fnmatch.translate(x) for x in file_types])
    files = [f for f in files if re.match(file_types, f)]
    basename_no_extension = os.path.splitext(os.path.basename(image_filename))[0]
    file_name_prefix = basename_no_extension + '.*'
    # files = [os.path.join(root, f) for f in files]
    files = [f for f in files if re.match(file_name_prefix, os.path.splitext(os.path.basename(f))[0])]
    return files
 
 
def main():
    # for root, _, files in os.walk(ANNOTATION_DIR):
    files = os.listdir(IMAGE_DIR)
    image_files = filter_for_pic(files)
    masks_generator(image_files)

    coco_output = {
        "info": INFO,
        "licenses": LICENSES,
        "categories": CATEGORIES,
        "images": [],
        "annotations": []
    }
    image_id = 1
    segmentation_id = 1
    
    files = os.listdir(INSTANCE_DIR)
    instance_files = filter_for_pic(files)

    # go through each image
    for image_filename in image_files:
        image_path = os.path.join(IMAGE_DIR, image_filename)
        image = Image.open(image_path)
        image_info = pycococreatortools.create_image_info(
                image_id, os.path.basename(image_filename), image.size)
        coco_output["images"].append(image_info)

        # filter for associated png annotations
        # for root, _, files in os.walk(INSTANCE_DIR):
        annotation_files = filter_for_instances(INSTANCE_DIR, instance_files, image_filename)

        # go through each associated annotation
        for annotation_filename in annotation_files:
            annotation_path = os.path.join(INSTANCE_DIR, annotation_filename)
            print(annotation_path)
            class_id = [x['id'] for x in CATEGORIES if x['name'] in annotation_filename][0]

            category_info = {'id': class_id, 'is_crowd': 'crowd' in image_filename}
            binary_mask = np.asarray(Image.open(annotation_path)
                                        .convert('1')).astype(np.uint8)

            annotation_info = pycococreatortools.create_annotation_info(
                    segmentation_id, image_id, category_info, binary_mask,
                    image.size, tolerance=2)

            if annotation_info is not None:
                coco_output["annotations"].append(annotation_info)

                segmentation_id = segmentation_id + 1

        image_id = image_id + 1
 
    with open('{}/frankfurt.json'.format(ROOT_DIR), 'w') as output_json_file:
        json.dump(coco_output, output_json_file)
 
 
if __name__ == "__main__":
    main()

代码的思路是根据cityscapes中的gtFine_instanceids.png获取每个instance的二值Mask图像并保存到INSTANCE_DIR中,文件名格式是:

          <原始图片文件名>_<类别名>_<实例序号>.png

随后,读入图片,并依次找出所有对应该图片的Mask,利用coco的API实现images和annotations的生成,并保存。

使用时,需要修改最开始的ROOT_DIR, IMAGES_DIR, ANNOTATION_DIR,并且在ROOT_DIR下创建存放实例Mask的INSTANCE_DIR。

 

此外,生成.json之后,最好查看以下标注的转换是否正确,博主也写了脚本,很短,大家稍加阅读就可以使用。

原始的COCO API中只能够像下图一样显示instance的polygon区域,没有bounding box以及class,不便于查看。

为此,在coco.showAnns基础上稍作了修改,加入绘制bbox以及左上角写类别,效果如下:

上述代码均已上传在Github:https://github.com/KevinJia1212/cityscapes2coco,如有疑问欢迎大家留言。

###############################################更新分割线####################################

应各位要求,我已将转换好的数据集上传至baidu网盘:

链接:https://pan.baidu.com/s/1bjDp9YsmXDx-TJk4R1jzhw 
提取码:au9b

注意:此数据集仅包含car、pedestrian(person)、truck、bus、rider五类标签,如需其他类别请自行利用转换代码转换。

  • 25
    点赞
  • 98
    收藏
    觉得还不错? 一键收藏
  • 47
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 47
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值