python提取COCO，VOC数据集中特定的类

柔情只为你懂 2021-11-22 13:54 324阅读 0赞

##### 1.python提取COCO数据集中特定的类 #####

安装pycocotools github地址：https://github.com/philferriere/cocoapi

pip install git+https://github.com/philferriere/cocoapi.git#subdirectory=PythonAPI

提取特定的类别如下：

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)

##### 2. 将上一步提取的COCO 某一类 xml转为COCO标准的json文件： #####

# -*- 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'))

##### 3.python提取Pascal Voc数据集中特定的类 #####

# -*- 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()