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

2024-01-25 19:46| 来源: 网络整理| 查看: 265

##￥##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五类标签，如需其他类别请自行利用转换代码转换。

【本文地址】

公司简介

联系我们