在網上看了很多將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/，
安裝一步到位，不用修改配置。
恩，大致就這樣，大家科研順利。