forked from andreasveit/coco-text
-
Notifications
You must be signed in to change notification settings - Fork 0
/
coco_text.py
234 lines (211 loc) · 9.98 KB
/
coco_text.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
__author__ = 'andreasveit'
__version__ = '1.1'
# Interface for accessing the COCO-Text dataset.
# COCO-Text is a large dataset designed for text detection and recognition.
# This is a Python API that assists in loading, parsing and visualizing the
# annotations. The format of the COCO-Text annotations is also described on
# the project website http://vision.cornell.edu/se3/coco-text/. In addition to this API, please download both
# the COCO images and annotations.
# This dataset is based on Microsoft COCO. Please visit http://mscoco.org/
# for more information on COCO, including for the image data, object annotatins
# and caption annotations.
# An alternative to using the API is to load the annotations directly
# into Python dictionary:
# with open(annotation_filename) as json_file:
# coco_text = json.load(json_file)
# Using the API provides additional utility functions.
# The following API functions are defined:
# COCO_Text - COCO-Text api class that loads COCO annotations and prepare data structures.
# getAnnIds - Get ann ids that satisfy given filter conditions.
# getImgIds - Get img ids that satisfy given filter conditions.
# loadAnns - Load anns with the specified ids.
# loadImgs - Load imgs with the specified ids.
# showAnns - Display the specified annotations.
# loadRes - Load algorithm results and create API for accessing them.
# Throughout the API "ann"=annotation, "cat"=category, and "img"=image.
# COCO-Text Toolbox. Version 1.1
# Data and paper available at: http://vision.cornell.edu/se3/coco-text/
# Code based on Microsoft COCO Toolbox Version 1.0 by Piotr Dollar and Tsung-Yi Lin
# extended and adapted by Andreas Veit, 2016.
# Licensed under the Simplified BSD License [see bsd.txt]
import json
import datetime
import matplotlib.pyplot as plt
from matplotlib.collections import PatchCollection
from matplotlib.patches import Rectangle, PathPatch
from matplotlib.path import Path
import numpy as np
import copy
import os
class COCO_Text:
def __init__(self, annotation_file=None):
"""
Constructor of COCO-Text helper class for reading and visualizing annotations.
:param annotation_file (str): location of annotation file
:return:
"""
# load dataset
self.dataset = {}
self.anns = {}
self.imgToAnns = {}
self.catToImgs = {}
self.imgs = {}
self.cats = {}
self.val = []
self.test = []
self.train = []
if not annotation_file == None:
assert os.path.isfile(annotation_file), "file does not exist"
print 'loading annotations into memory...'
time_t = datetime.datetime.utcnow()
dataset = json.load(open(annotation_file, 'r'))
print datetime.datetime.utcnow() - time_t
self.dataset = dataset
self.createIndex()
def createIndex(self):
# create index
print 'creating index...'
self.imgToAnns = {int(cocoid): self.dataset['imgToAnns'][cocoid] for cocoid in self.dataset['imgToAnns']}
self.imgs = {int(cocoid): self.dataset['imgs'][cocoid] for cocoid in self.dataset['imgs']}
self.anns = {int(annid): self.dataset['anns'][annid] for annid in self.dataset['anns']}
self.cats = self.dataset['cats']
self.val = [int(cocoid) for cocoid in self.dataset['imgs'] if self.dataset['imgs'][cocoid]['set'] == 'val']
self.test = [int(cocoid) for cocoid in self.dataset['imgs'] if self.dataset['imgs'][cocoid]['set'] == 'test']
self.train = [int(cocoid) for cocoid in self.dataset['imgs'] if self.dataset['imgs'][cocoid]['set'] == 'train']
print 'index created!'
def info(self):
"""
Print information about the annotation file.
:return:
"""
for key, value in self.dataset['info'].items():
print '%s: %s'%(key, value)
def filtering(self, filterDict, criteria):
return [key for key in filterDict if all(criterion(filterDict[key]) for criterion in criteria)]
def getAnnByCat(self, properties):
"""
Get ann ids that satisfy given properties
:param properties (list of tuples of the form [(category type, category)] e.g., [('readability','readable')]
: get anns for given categories - anns have to satisfy all given property tuples
:return: ids (int array) : integer array of ann ids
"""
return self.filtering(self.anns, [lambda d, x=a, y=b:d[x] == y for (a,b) in properties])
def getAnnIds(self, imgIds=[], catIds=[], areaRng=[]):
"""
Get ann ids that satisfy given filter conditions. default skips that filter
:param imgIds (int array) : get anns for given imgs
catIds (list of tuples of the form [(category type, category)] e.g., [('readability','readable')]
: get anns for given cats
areaRng (float array) : get anns for given area range (e.g. [0 inf])
:return: ids (int array) : integer array of ann ids
"""
imgIds = imgIds if type(imgIds) == list else [imgIds]
catIds = catIds if type(catIds) == list else [catIds]
if len(imgIds) == len(catIds) == len(areaRng) == 0:
anns = self.anns.keys()
else:
if not len(imgIds) == 0:
anns = sum([self.imgToAnns[imgId] for imgId in imgIds if imgId in self.imgToAnns],[])
else:
anns = self.anns.keys()
anns = anns if len(catIds) == 0 else list(set(anns).intersection(set(self.getAnnByCat(catIds))))
anns = anns if len(areaRng) == 0 else [ann for ann in anns if self.anns[ann]['area'] > areaRng[0] and self.anns[ann]['area'] < areaRng[1]]
return anns
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 type(imgIds) == list else [imgIds]
catIds = catIds if type(catIds) == list else [catIds]
if len(imgIds) == len(catIds) == 0:
ids = self.imgs.keys()
else:
ids = set(imgIds)
if not len(catIds) == 0:
ids = ids.intersection(set([self.anns[annid]['image_id'] for annid in self.getAnnByCat(catIds)]))
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 type(ids) == list:
return [self.anns[id] for id in ids]
elif type(ids) == int:
return [self.anns[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 type(ids) == list:
return [self.imgs[id] for id in ids]
elif type(ids) == int:
return [self.imgs[ids]]
def showAnns(self, anns, show_polygon=False):
"""
Display the specified annotations.
:param anns (array of object): annotations to display
:return: None
"""
if len(anns) == 0:
return 0
ax = plt.gca()
boxes = []
color = []
for ann in anns:
c = np.random.random((1, 3)).tolist()[0]
if show_polygon:
tl_x, tl_y, tr_x, tr_y, br_x, br_y, bl_x, bl_y = ann['polygon']
verts = [(tl_x, tl_y), (tr_x, tr_y), (br_x, br_y), (bl_x, bl_y), (0, 0)]
codes = [Path.MOVETO, Path.LINETO, Path.LINETO, Path.LINETO, Path.CLOSEPOLY]
path = Path(verts, codes)
patch = PathPatch(path, facecolor='none')
boxes.append(patch)
left, top = tl_x, tl_y
else:
left, top, width, height = ann['bbox']
boxes.append(Rectangle([left,top],width,height,alpha=0.4))
color.append(c)
if 'utf8_string' in ann.keys():
ax.annotate(ann['utf8_string'],(left,top-4),color=c)
p = PatchCollection(boxes, facecolors=color, edgecolors=(0,0,0,1), linewidths=3, alpha=0.4)
ax.add_collection(p)
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_Text()
res.dataset['imgs'] = [img for img in self.dataset['imgs']]
print 'Loading and preparing results... '
time_t = datetime.datetime.utcnow()
if type(resFile) == str:
anns = json.load(open(resFile))
else:
anns = resFile
assert type(anns) == list, 'results in not an array of objects'
annsImgIds = [int(ann['image_id']) for ann in anns]
if set(annsImgIds) != (set(annsImgIds) & set(self.getImgIds())):
print 'Results do not correspond to current coco set'
print 'skipping ', str(len(set(annsImgIds)) - len(set(annsImgIds) & set(self.getImgIds()))), ' images'
annsImgIds = list(set(annsImgIds) & set(self.getImgIds()))
res.imgToAnns = {cocoid : [] for cocoid in annsImgIds}
res.imgs = {cocoid: self.imgs[cocoid] for cocoid in annsImgIds}
assert anns[0]['bbox'] != [], 'results have incorrect format'
for id, ann in enumerate(anns):
if ann['image_id'] not in annsImgIds:
continue
bb = ann['bbox']
ann['area'] = bb[2]*bb[3]
ann['id'] = id
res.anns[id] = ann
res.imgToAnns[ann['image_id']].append(id)
print 'DONE (t=%0.2fs)'%((datetime.datetime.utcnow() - time_t).total_seconds())
return res