python如何实现提取COCO,VOC数据集中特定的类

from pycocotools.coco import COCO import os import shutil from tqdm import tqdm import skimage.io as io import matplotlib.pyplot as plt import cv2 from PIL import Image, ImageDraw #the path you want to save your results for coco to voc savepath="/media/huanglong/Newsmy/COCO/" #保存提取类的路径,我放在同一路径下 img_dir=savepath+'images/' anno_dir=savepath+'Annotations/' # datasets_list=['train2014', 'val2014'] datasets_list=['train2014'] classes_names = ['person'] #coco有80类，这里写要提取类的名字，以person为例 #Store annotations and train2014/val2014/... in this folder dataDir= '/media/huanglong/Newsmy/COCO/' #原coco数据集 headstr = """\ <annotation> <folder>VOC</folder> <filename>%s</filename> <source> <database>My Database</database> <annotation>COCO</annotation> <image>flickr</image> <flickrid>NULL</flickrid> </source> <owner> <flickrid>NULL</flickrid> <name>company</name> </owner> <size> <width>%d</width> <height>%d</height> <depth>%d</depth> </size> <segmented>0</segmented> """ objstr = """\ <object> <name>%s</name> <pose>Unspecified</pose> <truncated>0</truncated> <difficult>0</difficult> <bndbox> <xmin>%d</xmin> <ymin>%d</ymin> <xmax>%d</xmax> <ymax>%d</ymax> </bndbox> </object> """ tailstr = '''\ </annotation> ''' #if the dir is not exists,make it,else delete it def mkr(path): if os.path.exists(path): shutil.rmtree(path) os.mkdir(path) else: os.mkdir(path) mkr(img_dir) mkr(anno_dir) def id2name(coco): classes=dict() for cls in coco.dataset['categories']: classes[cls['id']]=cls['name'] return classes def write_xml(anno_path,head, objs, tail): f = open(anno_path, "w") f.write(head) for obj in objs: f.write(objstr%(obj[0],obj[1],obj[2],obj[3],obj[4])) f.write(tail) def save_annotations_and_imgs(coco,dataset,filename,objs): #eg:COCO_train2014_000000196610.jpg-->COCO_train2014_000000196610.xml anno_path=anno_dir+filename[:-3]+'xml' img_path=dataDir+dataset+'/'+filename print(img_path) dst_imgpath=img_dir+filename img=cv2.imread(img_path) #if (img.shape[2] == 1): # print(filename + " not a RGB image") # return shutil.copy(img_path, dst_imgpath) head=headstr % (filename, img.shape[1], img.shape[0], img.shape[2]) tail = tailstr write_xml(anno_path,head, objs, tail) def showimg(coco,dataset,img,classes,cls_id,show=True): global dataDir I=Image.open('%s/%s/%s'%(dataDir,dataset,img['file_name'])) #通过id，得到注释的信息 annIds = coco.getAnnIds(imgIds=img['id'], catIds=cls_id, iscrowd=None) # print(annIds) anns = coco.loadAnns(annIds) # print(anns) # coco.showAnns(anns) objs = [] for ann in anns: class_name=classes[ann['category_id']] if class_name in classes_names: print(class_name) if 'bbox' in ann: bbox=ann['bbox'] xmin = int(bbox[0]) ymin = int(bbox[1]) xmax = int(bbox[2] + bbox[0]) ymax = int(bbox[3] + bbox[1]) obj = [class_name, xmin, ymin, xmax, ymax] objs.append(obj) draw = ImageDraw.Draw(I) draw.rectangle([xmin, ymin, xmax, ymax]) if show: plt.figure() plt.axis('off') plt.imshow(I) plt.show() return objs for dataset in datasets_list: #./COCO/annotations/instances_train2014.json annFile='{}/annotations/instances_{}.json'.format(dataDir,dataset) #COCO API for initializing annotated data coco = COCO(annFile) #show all classes in coco classes = id2name(coco) print(classes) #[1, 2, 3, 4, 6, 8] classes_ids = coco.getCatIds(catNms=classes_names) print(classes_ids) for cls in classes_names: #Get ID number of this class cls_id=coco.getCatIds(catNms=[cls]) img_ids=coco.getImgIds(catIds=cls_id) print(cls,len(img_ids)) # imgIds=img_ids[0:10] for imgId in tqdm(img_ids): img = coco.loadImgs(imgId)[0] filename = img['file_name'] # print(filename) objs=showimg(coco, dataset, img, classes,classes_ids,show=False) print(objs) save_annotations_and_imgs(coco, dataset, filename, objs)

# -*- coding: utf-8 -*- # @Time : 2019/8/27 10：48 # @Author :Rock # @File : voc2coco.py # just for object detection import xml.etree.ElementTree as ET import os import json coco = dict() coco['images'] = [] coco['type'] = 'instances' coco['annotations'] = [] coco['categories'] = [] category_set = dict() image_set = set() category_item_id = 0 image_id = 0 annotation_id = 0 def addCatItem(name): global category_item_id category_item = dict() category_item['supercategory'] = 'none' category_item_id += 1 category_item['id'] = category_item_id category_item['name'] = name coco['categories'].append(category_item) category_set[name] = category_item_id return category_item_id def addImgItem(file_name, size): global image_id if file_name is None: raise Exception('Could not find filename tag in xml file.') if size['width'] is None: raise Exception('Could not find width tag in xml file.') if size['height'] is None: raise Exception('Could not find height tag in xml file.') img_id = "%04d" % image_id image_id += 1 image_item = dict() image_item['id'] = int(img_id) # image_item['id'] = image_id image_item['file_name'] = file_name image_item['width'] = size['width'] image_item['height'] = size['height'] coco['images'].append(image_item) image_set.add(file_name) return image_id def addAnnoItem(object_name, image_id, category_id, bbox): global annotation_id annotation_item = dict() annotation_item['segmentation'] = [] seg = [] # bbox[] is x,y,w,h # left_top seg.append(bbox[0]) seg.append(bbox[1]) # left_bottom seg.append(bbox[0]) seg.append(bbox[1] + bbox[3]) # right_bottom seg.append(bbox[0] + bbox[2]) seg.append(bbox[1] + bbox[3]) # right_top seg.append(bbox[0] + bbox[2]) seg.append(bbox[1]) annotation_item['segmentation'].append(seg) annotation_item['area'] = bbox[2] * bbox[3] annotation_item['iscrowd'] = 0 annotation_item['ignore'] = 0 annotation_item['image_id'] = image_id annotation_item['bbox'] = bbox annotation_item['category_id'] = category_id annotation_id += 1 annotation_item['id'] = annotation_id coco['annotations'].append(annotation_item) def parseXmlFiles(xml_path): for f in os.listdir(xml_path): if not f.endswith('.xml'): continue bndbox = dict() size = dict() current_image_id = None current_category_id = None file_name = None size['width'] = None size['height'] = None size['depth'] = None xml_file = os.path.join(xml_path, f) # print(xml_file) tree = ET.parse(xml_file) root = tree.getroot() if root.tag != 'annotation': raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag)) # elem is <folder>, <filename>, <size>, <object> for elem in root: current_parent = elem.tag current_sub = None object_name = None if elem.tag == 'folder': continue if elem.tag == 'filename': file_name = elem.text if file_name in category_set: raise Exception('file_name duplicated') # add img item only after parse <size> tag elif current_image_id is None and file_name is not None and size['width'] is not None: if file_name not in image_set: current_image_id = addImgItem(file_name, size) # print('add image with {} and {}'.format(file_name, size)) else: raise Exception('duplicated image: {}'.format(file_name)) # subelem is <width>, <height>, <depth>, <name>, <bndbox> for subelem in elem: bndbox['xmin'] = None bndbox['xmax'] = None bndbox['ymin'] = None bndbox['ymax'] = None current_sub = subelem.tag if current_parent == 'object' and subelem.tag == 'name': object_name = subelem.text if object_name not in category_set: current_category_id = addCatItem(object_name) else: current_category_id = category_set[object_name] elif current_parent == 'size': if size[subelem.tag] is not None: raise Exception('xml structure broken at size tag.') size[subelem.tag] = int(subelem.text) # option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox> for option in subelem: if current_sub == 'bndbox': if bndbox[option.tag] is not None: raise Exception('xml structure corrupted at bndbox tag.') bndbox[option.tag] = int(option.text) # only after parse the <object> tag if bndbox['xmin'] is not None: if object_name is None: raise Exception('xml structure broken at bndbox tag') if current_image_id is None: raise Exception('xml structure broken at bndbox tag') if current_category_id is None: raise Exception('xml structure broken at bndbox tag') bbox = [] # x bbox.append(bndbox['xmin']) # y bbox.append(bndbox['ymin']) # w bbox.append(bndbox['xmax'] - bndbox['xmin']) # h bbox.append(bndbox['ymax'] - bndbox['ymin']) # print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id, # bbox)) addAnnoItem(object_name, current_image_id, current_category_id, bbox) if __name__ == '__main__': #修改这里的两个地址，一个是xml文件的父目录；一个是生成的json文件的绝对路径 xml_path = r'G:\dataset\COCO\person\coco_val2014\annotations\\' json_file = r'G:\dataset\COCO\person\coco_val2014\instances_val2014.json' parseXmlFiles(xml_path) json.dump(coco, open(json_file, 'w'))

# -*- coding: utf-8 -*- # @Function:There are 20 classes in VOC data set. If you need to extract specific classes, you can use this program to extract them. import os import shutil ann_filepath='E:/VOCdevkit/VOC2012/Annotations/' img_filepath='E:/VOCdevkit/VOC2012/JPEGImages/' img_savepath='E:TrafficDatasets/JPEGImages/' ann_savepath='E:TrafficDatasets/Annotations/' if not os.path.exists(img_savepath): os.mkdir(img_savepath) if not os.path.exists(ann_savepath): os.mkdir(ann_savepath) names = locals() classes = ['aeroplane','bicycle','bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow','diningtable', 'dog', 'horse', 'motorbike', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor', 'person'] for file in os.listdir(ann_filepath): print(file) fp = open(ann_filepath + '\\' + file) #打开Annotations文件 ann_savefile=ann_savepath+file fp_w = open(ann_savefile, 'w') lines = fp.readlines() ind_start = [] ind_end = [] lines_id_start = lines[:] lines_id_end = lines[:] classes1 = '\t\t<name>bicycle</name>\n' classes2 = '\t\t<name>bus</name>\n' classes3 = '\t\t<name>car</name>\n' classes4 = '\t\t<name>motorbike</name>\n' classes5 = '\t\t<name>train</name>\n' #在xml中找到object块，并将其记录下来 while "\t<object>\n" in lines_id_start: a = lines_id_start.index("\t<object>\n") ind_start.append(a) #ind_start是<object>的行数 lines_id_start[a] = "delete" while "\t</object>\n" in lines_id_end: b = lines_id_end.index("\t</object>\n") ind_end.append(b) #ind_end是</object>的行数 lines_id_end[b] = "delete" #names中存放所有的object块 i = 0 for k in range(0, len(ind_start)): names['block%d' % k] = [] for j in range(0, len(classes)): if classes[j] in lines[ind_start[i] + 1]: a = ind_start[i] for o in range(ind_end[i] - ind_start[i] + 1): names['block%d' % k].append(lines[a + o]) break i += 1 #print(names['block%d' % k]) #xml头 string_start = lines[0:ind_start[0]] #xml尾 if((file[2:4]=='09') | (file[2:4]=='10') | (file[2:4]=='11')): string_end = lines[(len(lines) - 11):(len(lines))] else: string_end = [lines[len(lines) - 1]] #在给定的类中搜索，若存在则，写入object块信息 a = 0 for k in range(0, len(ind_start)): if classes1 in names['block%d' % k]: a += 1 string_start += names['block%d' % k] if classes2 in names['block%d' % k]: a += 1 string_start += names['block%d' % k] if classes3 in names['block%d' % k]: a += 1 string_start += names['block%d' % k] if classes4 in names['block%d' % k]: a += 1 string_start += names['block%d' % k] if classes5 in names['block%d' % k]: a += 1 string_start += names['block%d' % k] string_start += string_end # print(string_start) for c in range(0, len(string_start)): fp_w.write(string_start[c]) fp_w.close() #如果没有我们寻找的模块，则删除此xml，有的话拷贝图片 if a == 0: os.remove(ann_savepath+file) else: name_img = img_filepath + os.path.splitext(file)[0] + ".jpg" shutil.copy(name_img, img_savepath) fp.close()

相关阅读