devbox
2023面试高手
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

提取COCO 数据集的部分类

作者：2023面试高手 | 2024-05-25 16:01:51

赞

踩

提取COCO 数据集的部分类

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

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

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

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

换成

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

实在不行的话，手动下载

  1. git clone https://github.com/pdollar/coco.git
  2. cd coco/PythonAPI
  3. python setup.py build_ext --inplace #安装到本地
  4. python setup.py build_ext install # 安装到Python环境中

没有的库自己pip

注意skimage用pip install scikit-image -i https://pypi.tuna.tsinghua.edu.cn/simple

提取特定的类别如下：

  1. # conding='utf-8'
  2. from pycocotools.coco import COCO
  3. import os
  4. import shutil
  5. from tqdm import tqdm
  6. import skimage.io as io
  7. import matplotlib.pyplot as plt
  8. import cv2
  9. from PIL import Image, ImageDraw
  10.  
  11. #the path you want to save your results for coco to voc
  12. savepath="/opt/10T/home/asc005/YangMingxiang/DenseCLIP_/data/COCO/"  #save_path
  13. img_dir=savepath+'images/'
  14. anno_dir=savepath+'Annotations/'
  15. # datasets_list=['train2014', 'val2014']
  16. datasets_list=['train2017', 'val2017']
  17.  
  18. classes_names = ['sheep']  #coco
  19. #Store annotations and train2014/val2014/... in this folder
  20. dataDir= '/opt/10T/home/asc005/YangMingxiang/DenseCLIP_/data/coco/'  #origin coco
  21.  
  22. headstr = """\
  23. <annotation>
  24.     <folder>VOC</folder>
  25.     <filename>%s</filename>
  26.     <source>
  27.         <database>My Database</database>
  28.         <annotation>COCO</annotation>
  29.         <image>flickr</image>
  30.         <flickrid>NULL</flickrid>
  31.     </source>
  32.     <owner>
  33.         <flickrid>NULL</flickrid>
  34.         <name>company</name>
  35.     </owner>
  36.     <size>
  37.         <width>%d</width>
  38.         <height>%d</height>
  39.         <depth>%d</depth>
  40.     </size>
  41.     <segmented>0</segmented>
  42. """
  43. objstr = """\
  44.     <object>
  45.         <name>%s</name>
  46.         <pose>Unspecified</pose>
  47.         <truncated>0</truncated>
  48.         <difficult>0</difficult>
  49.         <bndbox>
  50.             <xmin>%d</xmin>
  51.             <ymin>%d</ymin>
  52.             <xmax>%d</xmax>
  53.             <ymax>%d</ymax>
  54.         </bndbox>
  55.     </object>
  56. """
  57.  
  58. tailstr = '''\
  59. </annotation>
  60. '''
  61.  
  62. #if the dir is not exists,make it,else delete it
  63. def mkr(path):
  64.     if os.path.exists(path):
  65.         shutil.rmtree(path)
  66.         os.mkdir(path)
  67.     else:
  68.         os.mkdir(path)
  69. mkr(img_dir)
  70. mkr(anno_dir)
  71. def id2name(coco):
  72.     classes=dict()
  73.     for cls in coco.dataset['categories']:
  74.         classes[cls['id']]=cls['name']
  75.     return classes
  76.  
  77. def write_xml(anno_path,head, objs, tail):
  78.     f = open(anno_path, "w")
  79.     f.write(head)
  80.     for obj in objs:
  81.         f.write(objstr%(obj[0],obj[1],obj[2],obj[3],obj[4]))
  82.     f.write(tail)
  83.  
  84.  
  85. def save_annotations_and_imgs(coco,dataset,filename,objs):
  86.     #eg:COCO_train2014_000000196610.jpg-->COCO_train2014_000000196610.xml
  87.     anno_path=anno_dir+filename[:-3]+'xml'
  88.     img_path=dataDir+dataset+'/'+filename
  89.     print(img_path)
  90.     dst_imgpath=img_dir+filename
  91.  
  92.     img=cv2.imread(img_path)
  93.     #if (img.shape[2] == 1):
  94.     #    print(filename + " not a RGB image")
  95.      #   return
  96.     shutil.copy(img_path, dst_imgpath)
  97.  
  98.     head=headstr % (filename, img.shape[1], img.shape[0], img.shape[2])
  99.     tail = tailstr
  100.     write_xml(anno_path,head, objs, tail)
  101.  
  102.  
  103. def showimg(coco,dataset,img,classes,cls_id,show=True):
  104.     global dataDir
  105.     I=Image.open('%s/%s/%s'%(dataDir,dataset,img['file_name']))
  106.     annIds = coco.getAnnIds(imgIds=img['id'], catIds=cls_id, iscrowd=None)
  107.     # print(annIds)
  108.     anns = coco.loadAnns(annIds)
  109.     # print(anns)
  110.     # coco.showAnns(anns)
  111.     objs = []
  112.     for ann in anns:
  113.         class_name=classes[ann['category_id']]
  114.         if class_name in classes_names:
  115.             print(class_name)
  116.             if 'bbox' in ann:
  117.                 bbox=ann['bbox']
  118.                 xmin = int(bbox[0])
  119.                 ymin = int(bbox[1])
  120.                 xmax = int(bbox[2] + bbox[0])
  121.                 ymax = int(bbox[3] + bbox[1])
  122.                 obj = [class_name, xmin, ymin, xmax, ymax]
  123.                 objs.append(obj)
  124.                 draw = ImageDraw.Draw(I)
  125.                 draw.rectangle([xmin, ymin, xmax, ymax])
  126.     if show:
  127.         plt.figure()
  128.         plt.axis('off')
  129.         plt.imshow(I)
  130.         plt.show()
  131.  
  132.     return objs
  133.  
  134. for dataset in datasets_list:
  135.     #./COCO/annotations/instances_train2014.json
  136.     annFile='{}/annotations/instances_{}.json'.format(dataDir,dataset)
  137.  
  138.     #COCO API for initializing annotated data
  139.     coco = COCO(annFile)
  140.  
  141.     #show all classes in coco
  142.     classes = id2name(coco)
  143.     print(classes)
  144.     #[1, 2, 3, 4, 6, 8]
  145.     classes_ids = coco.getCatIds(catNms=classes_names)
  146.     print(classes_ids)
  147.     for cls in classes_names:
  148.         #Get ID number of this class
  149.         cls_id=coco.getCatIds(catNms=[cls])
  150.         img_ids=coco.getImgIds(catIds=cls_id)
  151.         print(cls,len(img_ids))
  152.         # imgIds=img_ids[0:10]
  153.         for imgId in tqdm(img_ids):
  154.             img = coco.loadImgs(imgId)[0]
  155.             filename = img['file_name']
  156.             # print(filename)
  157.             objs=showimg(coco, dataset, img, classes,classes_ids,show=False)
  158.             print(objs)
  159.             save_annotations_and_imgs(coco, dataset, filename, objs)
  160.

然后就可以了

2. 将上面获取的数据集划分为训练集和测试集
  1. #conding='utf-8'
  2. import os
  3. import random
  4. from shutil import copy2
  5.  
  6. # origin
  7. image_original_path = "/opt/10T/home/asc005/YangMingxiang/DenseCLIP_/data/COCO/images"
  8. label_original_path = "/opt/10T/home/asc005/YangMingxiang/DenseCLIP_/data/COCO/Annotations"
  9.  
  10. # parent_path = os.path.dirname(os.getcwd())
  11. # parent_path = "D:\\AI_Find"
  12. # train_image_path = os.path.join(parent_path, "image_data/seed/train/images/")
  13. # train_label_path = os.path.join(parent_path, "image_data/seed/train/labels/")
  14. train_image_path = os.path.join("/opt/10T/home/asc005/YangMingxiang/DenseCLIP_/data/COCO/train2017")
  15. train_label_path = os.path.join("/opt/10T/home/asc005/YangMingxiang/DenseCLIP_/data/COCO/annotations/train2017")
  16. test_image_path = os.path.join("/opt/10T/home/asc005/YangMingxiang/DenseCLIP_/data/COCO/val2017")
  17. test_label_path = os.path.join("/opt/10T/home/asc005/YangMingxiang/DenseCLIP_/data/COCO/annotations/val2017")
  18.  
  19.  
  20. # test_image_path = os.path.join(parent_path, 'image_data/seed/val/images/')
  21. # test_label_path = os.path.join(parent_path, 'image_data/seed/val/labels/')
  22.  
  23.  
  24. def mkdir():
  25.     if not os.path.exists(train_image_path):
  26.         os.makedirs(train_image_path)
  27.     if not os.path.exists(train_label_path):
  28.         os.makedirs(train_label_path)
  29.  
  30.     if not os.path.exists(test_image_path):
  31.         os.makedirs(test_image_path)
  32.     if not os.path.exists(test_label_path):
  33.         os.makedirs(test_label_path)
  34.  
  35.  
  36. def main():
  37.     mkdir()
  38.     all_image = os.listdir(image_original_path)
  39.     for i in range(len(all_image)):
  40.         num = random.randint(1,5)
  41.         if num != 2:
  42.             copy2(os.path.join(image_original_path, all_image[i]), train_image_path)
  43.             train_index.append(i)
  44.         else:
  45.             copy2(os.path.join(image_original_path, all_image[i]), test_image_path)
  46.             val_index.append(i)
  47.  
  48.     all_label = os.listdir(label_original_path)
  49.     for i in train_index:
  50.             copy2(os.path.join(label_original_path, all_label[i]), train_label_path)
  51.     for i in val_index:
  52.             copy2(os.path.join(label_original_path, all_label[i]), test_label_path)
  53.  
  54.  
  55. if __name__ == '__main__':
  56.     train_index = []
  57.     val_index = []
  58.     main()
3.将上一步提取的COCO 某一类 xml转为COCO标准的json文件：
  1. # -*- coding: utf-8 -*-
  2. # @Time    : 2019/8/27 10：48
  3. # @Author  :Rock
  4. # @File    : voc2coco.py
  5. # just for object detection
  6. import xml.etree.ElementTree as ET
  7. import os
  8. import json
  9.  
  10. coco = dict()
  11. coco['images'] = []
  12. coco['type'] = 'instances'
  13. coco['annotations'] = []
  14. coco['categories'] = []
  15.  
  16. category_set = dict()
  17. image_set = set()
  18.  
  19. category_item_id = 0
  20. image_id = 0
  21. annotation_id = 0
  22.  
  23.  
  24. def addCatItem(name):
  25.     global category_item_id
  26.     category_item = dict()
  27.     category_item['supercategory'] = 'none'
  28.     category_item_id += 1
  29.     category_item['id'] = category_item_id
  30.     category_item['name'] = name
  31.     coco['categories'].append(category_item)
  32.     category_set[name] = category_item_id
  33.     return category_item_id
  34.  
  35.  
  36. def addImgItem(file_name, size):
  37.     global image_id
  38.     if file_name is None:
  39.         raise Exception('Could not find filename tag in xml file.')
  40.     if size['width'] is None:
  41.         raise Exception('Could not find width tag in xml file.')
  42.     if size['height'] is None:
  43.         raise Exception('Could not find height tag in xml file.')
  44.     img_id = "%04d" % image_id
  45.     image_id += 1
  46.     image_item = dict()
  47.     image_item['id'] = int(img_id)
  48.     # image_item['id'] = image_id
  49.     image_item['file_name'] = file_name
  50.     image_item['width'] = size['width']
  51.     image_item['height'] = size['height']
  52.     coco['images'].append(image_item)
  53.     image_set.add(file_name)
  54.     return image_id
  55.  
  56.  
  57. def addAnnoItem(object_name, image_id, category_id, bbox):
  58.     global annotation_id
  59.     annotation_item = dict()
  60.     annotation_item['segmentation'] = []
  61.     seg = []
  62.     # bbox[] is x,y,w,h
  63.     # left_top
  64.     seg.append(bbox[0])
  65.     seg.append(bbox[1])
  66.     # left_bottom
  67.     seg.append(bbox[0])
  68.     seg.append(bbox[1] + bbox[3])
  69.     # right_bottom
  70.     seg.append(bbox[0] + bbox[2])
  71.     seg.append(bbox[1] + bbox[3])
  72.     # right_top
  73.     seg.append(bbox[0] + bbox[2])
  74.     seg.append(bbox[1])
  75.  
  76.     annotation_item['segmentation'].append(seg)
  77.  
  78.     annotation_item['area'] = bbox[2] * bbox[3]
  79.     annotation_item['iscrowd'] = 0
  80.     annotation_item['ignore'] = 0
  81.     annotation_item['image_id'] = image_id
  82.     annotation_item['bbox'] = bbox
  83.     annotation_item['category_id'] = category_id
  84.     annotation_id += 1
  85.     annotation_item['id'] = annotation_id
  86.     coco['annotations'].append(annotation_item)
  87.  
  88.  
  89. def parseXmlFiles(xml_path):
  90.     for f in os.listdir(xml_path):
  91.         if not f.endswith('.xml'):
  92.             continue
  93.  
  94.         bndbox = dict()
  95.         size = dict()
  96.         current_image_id = None
  97.         current_category_id = None
  98.         file_name = None
  99.         size['width'] = None
  100.         size['height'] = None
  101.         size['depth'] = None
  102.  
  103.         xml_file = os.path.join(xml_path, f)
  104.         # print(xml_file)
  105.  
  106.         tree = ET.parse(xml_file)
  107.         root = tree.getroot()
  108.         if root.tag != 'annotation':
  109.             raise Exception('pascal voc xml root element should be annotation, rather than {}'.format(root.tag))
  110.  
  111.         # elem is <folder>, <filename>, <size>, <object>
  112.         for elem in root:
  113.             current_parent = elem.tag
  114.             current_sub = None
  115.             object_name = None
  116.  
  117.             if elem.tag == 'folder':
  118.                 continue
  119.  
  120.             if elem.tag == 'filename':
  121.                 file_name = elem.text
  122.                 if file_name in category_set:
  123.                     raise Exception('file_name duplicated')
  124.  
  125.             # add img item only after parse <size> tag
  126.             elif current_image_id is None and file_name is not None and size['width'] is not None:
  127.                 if file_name not in image_set:
  128.                     current_image_id = addImgItem(file_name, size)
  129.                     # print('add image with {} and {}'.format(file_name, size))
  130.                 else:
  131.                     raise Exception('duplicated image: {}'.format(file_name))
  132.                     # subelem is <width>, <height>, <depth>, <name>, <bndbox>
  133.             for subelem in elem:
  134.                 bndbox['xmin'] = None
  135.                 bndbox['xmax'] = None
  136.                 bndbox['ymin'] = None
  137.                 bndbox['ymax'] = None
  138.  
  139.                 current_sub = subelem.tag
  140.                 if current_parent == 'object' and subelem.tag == 'name':
  141.                     object_name = subelem.text
  142.                     if object_name not in category_set:
  143.                         current_category_id = addCatItem(object_name)
  144.                     else:
  145.                         current_category_id = category_set[object_name]
  146.  
  147.                 elif current_parent == 'size':
  148.                     if size[subelem.tag] is not None:
  149.                         raise Exception('xml structure broken at size tag.')
  150.                     size[subelem.tag] = int(subelem.text)
  151.  
  152.                 # option is <xmin>, <ymin>, <xmax>, <ymax>, when subelem is <bndbox>
  153.                 for option in subelem:
  154.                     if current_sub == 'bndbox':
  155.                         if bndbox[option.tag] is not None:
  156.                             raise Exception('xml structure corrupted at bndbox tag.')
  157.                         bndbox[option.tag] = int(option.text)
  158.  
  159.                 # only after parse the <object> tag
  160.                 if bndbox['xmin'] is not None:
  161.                     if object_name is None:
  162.                         raise Exception('xml structure broken at bndbox tag')
  163.                     if current_image_id is None:
  164.                         raise Exception('xml structure broken at bndbox tag')
  165.                     if current_category_id is None:
  166.                         raise Exception('xml structure broken at bndbox tag')
  167.                     bbox = []
  168.                     # x
  169.                     bbox.append(bndbox['xmin'])
  170.                     # y
  171.                     bbox.append(bndbox['ymin'])
  172.                     # w
  173.                     bbox.append(bndbox['xmax'] - bndbox['xmin'])
  174.                     # h
  175.                     bbox.append(bndbox['ymax'] - bndbox['ymin'])
  176.                     # print('add annotation with {},{},{},{}'.format(object_name, current_image_id, current_category_id,
  177.                     #                                                bbox))
  178.                     addAnnoItem(object_name, current_image_id, current_category_id, bbox)
  179.  
  180.  
  181. if __name__ == '__main__':
  182. 	#修改这里的两个地址，一个是xml文件的父目录；一个是生成的json文件的绝对路径
  183.     xml_path = r'G:\dataset\COCO\person\coco_val2014\annotations\\'
  184.     json_file = r'G:\dataset\COCO\person\coco_val2014\instances_val2014.json'
  185.     parseXmlFiles(xml_path)
  186.     json.dump(coco, open(json_file, 'w'))
  187.

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/2023面试高手/article/detail/622710
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号