在网上看了很多将coco数据集进行类别提取的代码，但是都只是简单的将bbox数据进行提取而不提取segmentation数据，以前没接触过COCO数据集的我只能硬啃COCO的数据结构，结合前人的代码谢了一份提取完整特征的代码。
PS：这套代码虽然能提取部分类别的信息，但是我没有将原有类别的编号进行重新排序，如我在提取car，bus，truck三类信息时，其原来的类别编号是3，6，8，我按照原本的编号进行保存了，并没有重新排列成1，2，3或0，1，2进行保存，如果大家因为提取类别较多想重新排列类别，有两种方法：
第一种方法：需要将categories下的id字段和annotations下的category_id字段对应的进行更改;
第二种方法：不改变这套代码，而是在之后的数据处理中建一个label_map直接进行映射，方便快捷。(可以参考github中Yolact++算法代码下的data/config.py中的配置)
1 # -*- coding: utf-8 -*- 7 from pycocotools.coco import COCO 8 import os 9 import shutil 10 from tqdm import tqdm 11 import matplotlib.pyplot as plt 12 import cv2 13 from PIL import Image, ImageDraw 14 import skimage.io as io 15 import json 16 17 ''' 18 路径参数 19 ''' 20 #原coco数据集的路径 21 dataDir= './' 22 #用于保存新生成的数据的路径 23 savepath = "newdata/" 24 #只保存含有你需要的类别的图片的路径，最后没有用 25 #因为没必要，coco是按json中的信息读图，只要在json里做筛选就行了 26 img_save = savepath + 'images/' 27 #最后生产的json文件的保存路径 28 anno_save = savepath+'annotations/' 29 ''' 30 数据集参数 31 ''' 32 #coco有80类，这里写要提取部分类的名字 33 #如我只需要car、bus、truck这三类数据 34 classes_names = ['car','bus','truck'] 35 #要处理的数据集，比如val2017、train2017等 36 #不建议多个数据集在一个list中，我就跑崩了 37 #还是一次提取一个数据集安全点_(:3」∠❀)_ 38 datasets_list=['val2017'] 39 40 #生成保存路径，函数抄的(›´ω`‹ ) 41 #if the dir is not exists,make it,else delete it 42 def mkr(path): 43 if os.path.exists(path): 44 shutil.rmtree(path) 45 os.mkdir(path) 46 else: 47 os.mkdir(path) 48 49 #获取并处理所有需要的json数据 50 def process_json_data(annFile): 51 #获取COCO_json的数据 52 coco = COCO(annFile) 53 #拿到所有需要的图片数据的id 54 classes_ids = coco.getCatIds(catNms = classes_names) 55 #加载所有需要的类别信息 56 classes_list = coco.loadCats(classes_ids) 57 #取所有类别的并集的所有图片id 58 #如果想要交集，不需要循环，直接把所有类别作为参数输入，即可得到所有类别都包含的图片 59 imgIds_list = [] 60 for idx in classes_ids: 61 imgidx = coco.getImgIds(catIds=idx) 62 imgIds_list += imgidx 63 #去除重复的图片 64 imgIds_list = list(set(imgIds_list)) 65 #一次性获取所有图像的信息 66 image_info_list = coco.loadImgs(imgIds_list) 67 #获取图像中对应类别的分割信息,由catIds来指定 68 annIds = coco.getAnnIds(imgIds = [], catIds = classes_ids, iscrowd=None) 69 anns_list = coco.loadAnns(annIds) 70 return classes_list,image_info_list,anns_list 71 72 #保存数据到json 73 def save_json_data(json_file,classes_list,image_info_list,anns_list): 74 coco_sub = dict() 75 coco_sub['info'] = dict() 76 coco_sub['licenses'] = [] 77 coco_sub['images'] = [] 78 coco_sub['type'] = 'instances' 79 coco_sub['annotations'] = [] 80 coco_sub['categories'] = [] 81 #以下非必须,为coco数据集的前缀信息 82 coco_sub['info']['description'] = 'COCO 2017 sub Dataset' 83 coco_sub['info']['url'] = 'https://www.cnblogs.com/lhdb/' 84 coco_sub['info']['version'] = '1.0' 85 coco_sub['info']['year'] = 2020 86 coco_sub['info']['contributor'] = 'smh' 87 coco_sub['info']['date_created'] = '2020-7-1 10:06' 88 sub_license = dict() 89 sub_license['url'] = 'https://www.cnblogs.com/lhdb/' 90 sub_license['id'] = 1 91 sub_license['name'] = 'Attribution-NonCommercial-ShareAlike License' 92 coco_sub['licenses'].append(sub_license) 93 #以下为必须插入信息,包括image、annotations、categories三个字段 94 #插入image信息 95 coco_sub['images'].extend(image_info_list) 96 #插入annotation信息 97 coco_sub['annotations'].extend(anns_list) 98 #插入categories信息 99 coco_sub['categories'].extend(classes_list) 100 #自此所有该插入的数据就已经插入完毕啦٩( ๑╹ ꇴ╹)۶ 101 #最后一步，保存数据 102 json.dump(coco_sub, open(json_file, 'w')) 103 104 105 if __name__ == '__main__': 106 mkr(img_save) 107 mkr(anno_save) 108 #按单个数据集进行处理 109 for dataset in datasets_list: 110 #获取要处理的json文件路径 111 annFile='{}/annotations/instances_{}.json'.format(dataDir,dataset) 112 #存储处理完成的json文件路径 113 json_file = '{}/instances_{}_sub.json'.format(anno_save,dataset) 114 #处理数据 115 classes_list,image_info_list,anns_list = process_json_data(annFile) 116 #保存数据 117 save_json_data(json_file,classes_list,image_info_list,anns_list) 118 print('instances_{}_sub.json saved ٩( ๑╹ ꇴ ╹)۶'.format(dataset))
对了，大家在windows下安装pycocotools建议使用https://pypi.org/project/pycocotools-windows/，
安装一步到位，不用修改配置。
恩，大致就这样，大家科研顺利。