下面的代码改写自 COCO 官方 API，主要是替换 defaultdict 为 dict.get():
import time
import json
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import PatchCollection
from matplotlib.patches import Polygon
import copy
from urllib.request import urlretrieve
import itertools
class Bunch(dict):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.__dict__ = self
def _isArrayLike(obj):
'''
用于判断对象是否包含对应的属性
'''
return hasattr(obj, '__iter__') and hasattr(obj, '__len__')
class COCO(dict):
def __init__(self, annotation_file=None, *args, **kwargs):
"""
Constructor of Microsoft COCO helper class for reading and visualizing annotations.
:param annotation_file (str): location of annotation file
:param image_folder (str): location to the folder that hosts images.
:return:
"""
super().__init__(*args, **kwargs)
self.__dict__ = self
# load dataset
if annotation_file:
print('loading annotations into memory...')
tic = time.time()
dataset = json.load(open(annotation_file, 'r'))
assert type(
dataset
) == dict, 'annotation file format {} not supported'.format(
type(dataset))
print('Done (t={:0.2f}s)'.format(time.time() - tic))
self.dataset = Bunch(dataset)
self.createIndex()
def createIndex(self):
'''
创建索引
'''
print('creating index...')
imgToAnns, catToImgs = {}, {}
anns, cats, imgs = {}, {}, {}
if 'annotations' in self.dataset:
for ann in self.dataset['annotations']:
imgToAnns[ann['image_id']] = imgToAnns.get(
ann['image_id'], []) + [ann]
anns[ann['id']] = ann
if 'images' in self.dataset:
for img in self.dataset['images']:
imgs[img['id']] = img
if 'categories' in self.dataset:
for cat in self.dataset['categories']:
cats[cat['id']] = cat
if 'annotations' and 'categories' in self.dataset:
for ann in self.dataset['annotations']:
catToImgs[ann['category_id']] = catToImgs.get(
ann['category_id'], []) + [ann['image_id']]
print('index created!')
# create class members
self.anns = anns
self.imgToAnns = imgToAnns
self.catToImgs = catToImgs
self.imgs = imgs
self.cats = cats
def __str__(self):
"""
Print information about the annotation file.
:return:
"""
return ''.join('{}: {}\n'.format(key, value)
for key, value in coco.dataset['info'].items())
def getAnnIds(self, imgIds=[], catIds=[], areaRng=[], iscrowd=None):
"""
Get ann ids that satisfy given filter conditions. default skips that filter
:param imgIds (int array) : get anns for given imgs
catIds (int array) : get anns for given cats
areaRng (float array) : get anns for given area range (e.g. [0 inf])
iscrowd (boolean) : get anns for given crowd label (False or True)
:return: ids (int array) : integer array of ann ids
"""
imgIds = imgIds if _isArrayLike(imgIds) else [imgIds]
catIds = catIds if _isArrayLike(catIds) else [catIds]
if len(imgIds) == len(catIds) == len(areaRng) == 0:
anns = self.dataset['annotations']
else:
if not len(imgIds) == 0:
lists = [
self.imgToAnns[imgId] for imgId in imgIds
if imgId in self.imgToAnns
]
anns = list(itertools.chain.from_iterable(lists))
else:
anns = self.dataset['annotations']
anns = anns if len(catIds) == 0 else [
ann for ann in anns if ann['category_id'] in catIds
]
anns = anns if len(areaRng) == 0 else [
ann for ann in anns
if ann['area'] > areaRng[0] and ann['area'] < areaRng[1]
]
if not iscrowd == None:
ids = [ann['id'] for ann in anns if ann['iscrowd'] == iscrowd]
else:
ids = [ann['id'] for ann in anns]
return ids
def getCatIds(self, catNms=[], supNms=[], catIds=[]):
"""
filtering parameters. default skips that filter.
:param catNms (str array) : get cats for given cat names
:param supNms (str array) : get cats for given supercategory names
:param catIds (int array) : get cats for given cat ids
:return: ids (int array) : integer array of cat ids
"""
catNms = catNms if _isArrayLike(catNms) else [catNms]
supNms = supNms if _isArrayLike(supNms) else [supNms]
catIds = catIds if _isArrayLike(catIds) else [catIds]
if len(catNms) == len(supNms) == len(catIds) == 0:
cats = self.dataset['categories']
else:
cats = self.dataset['categories']
cats = cats if len(catNms) == 0 else [
cat for cat in cats if cat['name'] in catNms
]
cats = cats if len(supNms) == 0 else [
cat for cat in cats if cat['supercategory'] in supNms
]
cats = cats if len(catIds) == 0 else [
cat for cat in cats if cat['id'] in catIds
]
ids = [cat['id'] for cat in cats]
return ids
def getImgIds(self, imgIds=[], catIds=[]):
'''
Get img ids that satisfy given filter conditions.
:param imgIds (int array) : get imgs for given ids
:param catIds (int array) : get imgs with all given cats
:return: ids (int array) : integer array of img ids
'''
imgIds = imgIds if _isArrayLike(imgIds) else [imgIds]
catIds = catIds if _isArrayLike(catIds) else [catIds]
if len(imgIds) == len(catIds) == 0:
ids = self.imgs.keys()
else:
ids = set(imgIds)
for i, catId in enumerate(catIds):
if i == 0 and len(ids) == 0:
ids = set(self.catToImgs[catId])
else:
ids &= set(self.catToImgs[catId])
return list(ids)
def loadAnns(self, ids=[]):
"""
Load anns with the specified ids.
:param ids (int array) : integer ids specifying anns
:return: anns (object array) : loaded ann objects
"""
if _isArrayLike(ids):
return [self.anns[id] for id in ids]
elif type(ids) == int:
return [self.anns[ids]]
def loadCats(self, ids=[]):
"""
Load cats with the specified ids.
:param ids (int array) : integer ids specifying cats
:return: cats (object array) : loaded cat objects
"""
if _isArrayLike(ids):
return [self.cats[id] for id in ids]
elif type(ids) == int:
return [self.cats[ids]]
def loadImgs(self, ids=[]):
"""
Load anns with the specified ids.
:param ids (int array) : integer ids specifying img
:return: imgs (object array) : loaded img objects
"""
if _isArrayLike(ids):
return [self.imgs[id] for id in ids]
elif type(ids) == int:
return [self.imgs[ids]]
def showAnns(self, anns):
"""
Display the specified annotations.
:param anns (array of object): annotations to display
:return: None
"""
if len(anns) == 0:
return 0
if 'segmentation' in anns[0] or 'keypoints' in anns[0]:
datasetType = 'instances'
elif 'caption' in anns[0]:
datasetType = 'captions'
else:
raise Exception('datasetType not supported')
if datasetType == 'instances':
ax = plt.gca()
ax.set_autoscale_on(False)
polygons = []
color = []
for ann in anns:
c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0]
if 'segmentation' in ann:
if type(ann['segmentation']) == list:
# polygon
for seg in ann['segmentation']:
poly = np.array(seg).reshape((int(len(seg) / 2),
2))
polygons.append(Polygon(poly))
color.append(c)
else:
# mask
t = self.imgs[ann['image_id']]
if type(ann['segmentation']['counts']) == list:
rle = maskUtils.frPyObjects(
[ann['segmentation']], t['height'], t['width'])
else:
rle = [ann['segmentation']]
m = maskUtils.decode(rle)
img = np.ones((m.shape[0], m.shape[1], 3))
if ann['iscrowd'] == 1:
color_mask = np.array([2.0, 166.0, 101.0]) / 255
if ann['iscrowd'] == 0:
color_mask = np.random.random((1, 3)).tolist()[0]
for i in range(3):
img[:, :, i] = color_mask[i]
ax.imshow(np.dstack((img, m * 0.5)))
if 'keypoints' in ann and type(ann['keypoints']) == list:
# turn skeleton into zero-based index
sks = np.array(
self.loadCats(ann['category_id'])[0]['skeleton']) - 1
kp = np.array(ann['keypoints'])
x = kp[0::3]
y = kp[1::3]
v = kp[2::3]
for sk in sks:
if np.all(v[sk] > 0):
plt.plot(x[sk], y[sk], linewidth=3, color=c)
plt.plot(
x[v > 0],
y[v > 0],
'o',
markersize=8,
markerfacecolor=c,
markeredgecolor='k',
markeredgewidth=2)
plt.plot(
x[v > 1],
y[v > 1],
'o',
markersize=8,
markerfacecolor=c,
markeredgecolor=c,
markeredgewidth=2)
p = PatchCollection(
polygons, facecolor=color, linewidths=0, alpha=0.4)
ax.add_collection(p)
p = PatchCollection(
polygons, facecolor='none', edgecolors=color, linewidths=2)
ax.add_collection(p)
elif datasetType == 'captions':
for ann in anns:
print(ann['caption'])
def download(self, tarDir=None, imgIds=[]):
'''
Download COCO images from mscoco.org server.
:param tarDir (str): COCO results directory name
imgIds (list): images to be downloaded
:return:
'''
if tarDir is None:
print('Please specify target directory')
return -1
if len(imgIds) == 0:
imgs = self.imgs.values()
else:
imgs = self.loadImgs(imgIds)
N = len(imgs)
if not os.path.exists(tarDir):
os.makedirs(tarDir)
for i, img in enumerate(imgs):
tic = time.time()
fname = os.path.join(tarDir, img['file_name'])
if not os.path.exists(fname):
urlretrieve(img['coco_url'], fname)
print('downloaded {}/{} images (t={:0.1f}s)'.format(
i, N,
time.time() - tic))dataDir = 'D:datasets/coco/'
dataType = 'val2017'
annotation_file = '{}annotations/instances_{}.json'.format(dataDir, dataType)初始化 API
# initialize COCO api for instance annotations
coco = COCO(annotation_file)loading annotations into memory...
Done (t=0.80s)
creating index...
index created!如果你需要预览你载入的 COCO 数据集，可以使用 print() 来实现：
print(coco)description: COCO 2017 Dataset
url: http://cocodataset.org
version: 1.0
year: 2017
contributor: COCO Consortium
date_created: 2017/09/01coco.keys()dict_keys(['dataset', 'anns', 'imgToAnns', 'catToImgs', 'imgs', 'cats'])展示 COCO 的类别与超类
cats = coco.loadCats(coco.getCatIds())
nms = set([cat['name'] for cat in cats]) # 获取 cat 的 name 信息
print('COCO categories: \n{}\n'.format(' '.join(nms)))
# ============================================================
snms = set([cat['supercategory'] for cat in cats]) # 获取 cat 的 name 信息
print('COCO supercategories: \n{}'.format(' '.join(snms)))COCO categories:
broccoli motorcycle bench umbrella sheep skateboard cake cup mouse bus microwave baseball glove airplane person backpack bear giraffe hot dog couch orange dog truck dining table stop sign bottle cell phone oven toaster bird toothbrush handbag cat train traffic light cow zebra donut bed refrigerator boat pizza surfboard fork laptop knife banana clock teddy bear potted plant toilet snowboard wine glass bowl car tv kite carrot sink scissors sports ball sandwich frisbee spoon chair tennis racket horse apple bicycle tie baseball bat suitcase skis keyboard vase hair drier parking meter elephant remote book fire hydrant
COCO supercategories:
vehicle indoor accessory kitchen food sports furniture electronic person animal outdoor appliance通过给定条件获取图片
获取包含给定类别的所有图片
# get all images containing given categories, select one at random
catIds = coco.getCatIds(catNms=['cat', 'dog', 'snowboar']) # 获取 Cat 的 Ids
imgIds = coco.getImgIds(catIds=catIds ) #
img = coco.loadImgs(imgIds)随机选择一张图片的信息：
img = coco.loadImgs(imgIds[np.random.randint(0,len(imgIds))])[0]
img{'license': 1,
'file_name': '000000022892.jpg',
'coco_url': 'http://images.cocodataset.org/val2017/000000022892.jpg',
'height': 334,
'width': 500,
'date_captured': '2013-11-18 15:11:22',
'flickr_url': 'http://farm1.staticflickr.com/243/521936273_d0817d38a4_z.jpg',
'id': 22892}从网页获取图片
import skimage.io as sio%pylab inline
# use url to load image
I = sio.imread(img['coco_url'])
plt.axis('off')
plt.imshow(I)
plt.show()Populating the interactive namespace from numpy and matplotlib从本地读取图片
为了避免解压数据集，我使用了 zipfile 模块：
import zipfile
import cv2val_path = 'E:/Data/coco/val2017.zip'
val_z = zipfile.ZipFile(val_path)for name in val_z.namelist():
if name.endswith(img['file_name']):
bI = val_z.read(name)
I = cv2.imdecode(np.frombuffer(bI, np.uint8), cv2.IMREAD_ANYCOLOR)print('图片的尺寸：', I.shape)图片的尺寸： (334, 500, 3)plt.imshow(I)
plt.axis('off')
plt.show()将图片的 anns 信息标注在图片上
# load and display instance annotations
plt.imshow(I)
plt.axis('off')
annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
anns = coco.loadAnns(annIds)
coco.showAnns(anns)这里有一个梗：cv2 的图片默认模式是 BGR 而不是 RGB,所以，将 I 直接使用 plt 会改变原图的颜色空间，为此我们可以使用 cv2.COLOR_BGR2RGB.
I_ = cv2.cvtColor(I, cv2.COLOR_BGR2RGB)
# load and display instance annotations
plt.imshow(I_)
plt.axis('off')
annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
anns = coco.loadAnns(annIds)
coco.showAnns(anns)关键点检测
# initialize COCO api for person keypoints annotations
annFile = '{}/annotations/person_keypoints_{}.json'.format(dataDir, dataType)
coco_kps=COCO(annFile)loading annotations into memory...
Done (t=0.47s)
creating index...
index created!先选择一张带有 person 的图片：
catIds = coco.getCatIds(catNms=['person']) # 获取 Cat 的 Ids
imgIds = coco.getImgIds(catIds=catIds) #
img = coco.loadImgs(imgIds)[7]# use url to load image
I = sio.imread(img['coco_url'])
plt.axis('off')
plt.imshow(I)
plt.show()# load and display keypoints annotations
plt.imshow(I); plt.axis('off')
ax = plt.gca()
annIds = coco_kps.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
anns = coco_kps.loadAnns(annIds)
coco_kps.showAnns(anns)看图说话
# initialize COCO api for caption annotations
annFile = '{}/annotations/captions_{}.json'.format(dataDir,dataType)
coco_caps=COCO(annFile)loading annotations into memory...
Done (t=0.06s)
creating index...
index created!# load and display caption annotations
annIds = coco_caps.getAnnIds(imgIds=img['id']);
anns = coco_caps.loadAnns(annIds)
coco_caps.showAnns(anns)
plt.imshow(I); plt.axis('off'); plt.show()A woman is playing tennis with her children.
A mother and her children play on a grass tennis court
A woman hitting a tennis ball with a racquet
The woman serves the tennis ball as a child watches.
A woman hits a tennis ball with some kids.测试
这里还没有改好，但是应该足够了！
class RecCOCO(COCO):
def __init__(self, annotation_file=None, *args, **kwargs):
super().__init__(*args, **kwargs)
def loadRes(self, resFile):
"""
Load result file and return a result api object.
:param resFile (str) : file name of result file
:return: res (obj) : result api object
"""
res = COCO()
res.dataset['images'] = [img for img in self.dataset['images']]
print('Loading and preparing results...')
tic = time.time()
if type(resFile) == str or type(resFile) == unicode:
anns = json.load(open(resFile))
elif type(resFile) == np.ndarray:
anns = self.loadNumpyAnnotations(resFile)
else:
anns = resFile
assert type(anns) == list, 'results in not an array of objects'
annsImgIds = [ann['image_id'] for ann in anns]
assert set(annsImgIds) == (set(annsImgIds) & set(self.getImgIds())), \
'Results do not correspond to current coco set'
if 'caption' in anns[0]:
imgIds = set([img['id'] for img in res.dataset['images']]) & set(
[ann['image_id'] for ann in anns])
res.dataset['images'] = [
img for img in res.dataset['images'] if img['id'] in imgIds
]
for id, ann in enumerate(anns):
ann['id'] = id + 1
elif 'bbox' in anns[0] and not anns[0]['bbox'] == []:
res.dataset['categories'] = copy.deepcopy(
self.dataset['categories'])
for id, ann in enumerate(anns):
bb = ann['bbox']
x1, x2, y1, y2 = [bb[0], bb[0] + bb[2], bb[1], bb[1] + bb[3]]
if not 'segmentation' in ann:
ann['segmentation'] = [[x1, y1, x1, y2, x2, y2, x2, y1]]
ann['area'] = bb[2] * bb[3]
ann['id'] = id + 1
ann['iscrowd'] = 0
elif 'segmentation' in anns[0]:
res.dataset['categories'] = copy.deepcopy(
self.dataset['categories'])
for id, ann in enumerate(anns):
# now only support compressed RLE format as segmentation results
ann['area'] = maskUtils.area(ann['segmentation'])
if not 'bbox' in ann:
ann['bbox'] = maskUtils.toBbox(ann['segmentation'])
ann['id'] = id + 1
ann['iscrowd'] = 0
elif 'keypoints' in anns[0]:
res.dataset['categories'] = copy.deepcopy(
self.dataset['categories'])
for id, ann in enumerate(anns):
s = ann['keypoints']
x = s[0::3]
y = s[1::3]
x0, x1, y0, y1 = np.min(x), np.max(x), np.min(y), np.max(y)
ann['area'] = (x1 - x0) * (y1 - y0)
ann['id'] = id + 1
ann['bbox'] = [x0, y0, x1 - x0, y1 - y0]
print('DONE (t={:0.2f}s)'.format(time.time() - tic))
res.dataset['annotations'] = anns
res.createIndex()
return res
def loadNumpyAnnotations(self, data):
"""
Convert result data from a numpy array [Nx7] where each row contains {imageID,x1,y1,w,h,score,class}
:param data (numpy.ndarray)
:return: annotations (python nested list)
"""
print('Converting ndarray to lists...')
assert (type(data) == np.ndarray)
print(data.shape)
assert (data.shape[1] == 7)
N = data.shape[0]
ann = []
for i in range(N):
if i % 1000000 == 0:
print('{}/{}'.format(i, N))
ann += [{
'image_id': int(data[i, 0]),
'bbox': [data[i, 1], data[i, 2], data[i, 3], data[i, 4]],
'score': data[i, 5],
'category_id': int(data[i, 6]),
}]
return ann
def annToRLE(self, ann):
"""
Convert annotation which can be polygons, uncompressed RLE to RLE.
:return: binary mask (numpy 2D array)
"""
t = self.imgs[ann['image_id']]
h, w = t['height'], t['width']
segm = ann['segmentation']
if type(segm) == list:
# polygon -- a single object might consist of multiple parts
# we merge all parts into one mask rle code
rles = maskUtils.frPyObjects(segm, h, w)
rle = maskUtils.merge(rles)
elif type(segm['counts']) == list:
# uncompressed RLE
rle = maskUtils.frPyObjects(segm, h, w)
else:
# rle
rle = ann['segmentation']
return rle
def annToMask(self, ann):
"""
Convert annotation which can be polygons, uncompressed RLE, or RLE to binary mask.
:return: binary mask (numpy 2D array)
"""
rle = self.annToRLE(ann)
m = maskUtils.decode(rle)
return m