python如何实现提取COCO,VOC数据集中特定的类-创新互联

这篇文章主要介绍了python如何实现提取COCO,VOC数据集中特定的类，具有一定借鉴价值，感兴趣的朋友可以参考下，希望大家阅读完这篇文章之后大有收获，下面让小编带着大家一起了解一下。

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 = """\

 VOC
 %s
 
 My Database
 COCO
 flickr
 NULL
 
 
 NULL
 company
 
 
 %d
 %d
 %d
 
 0
"""
objstr = """\
 
"""
 
tailstr = '''\

'''
 
#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 , , ,