Skip to content

COCO module

COCOAnnotation

Bases: COCOBase

see https://cocodataset.org/#format-data

see https://github.com/cocodataset/cocoapi/issues/184

annotation {
    "id": int,
    "image_id": int,
    "category_id": int,
    "segmentation": RLE or [polygon],
    "area": float,
    "bbox": [x,y,width,height],
    "iscrowd": 0 or 1,
    "keypoints": [x1,y1,v1...],
    "num_keypoints": int
}

keypoints should be 3*len(keypoints_category), where keypoints_category is the keypoints in the corresponding category

Source code in pycocowriter/coco.py 
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
class COCOAnnotation(COCOBase):
    '''
        see https://cocodataset.org/#format-data

        see https://github.com/cocodataset/cocoapi/issues/184

            annotation {
                "id": int,
                "image_id": int,
                "category_id": int,
                "segmentation": RLE or [polygon],
                "area": float,
                "bbox": [x,y,width,height],
                "iscrowd": 0 or 1,
                "keypoints": [x1,y1,v1...],
                "num_keypoints": int
            }

        keypoints should be 3*len(keypoints_category), where keypoints_category is the keypoints in the corresponding category
    '''
    def __init__(self, image_id:int, eye_d:int, category_id:int, bbox:tuple[int]=None, area:float=None, segmentation=None, iscrowd:int=None, keypoints:list[int]=None):
        self.image_id = image_id
        self.id = eye_d
        self.category_id = category_id
        self.bbox = bbox
        area = area or self._compute_area()
        self.area = area
        self.segmentation = segmentation
        self.iscrowd = iscrowd or 0
        self.keypoints = keypoints
        self.num_keypoints = self._compute_num_keypoints()

    def _compute_num_keypoints(self):
        if self.keypoints is not None:
            # the number of keypoints is the number of "visible" keypoints.
            return sum([v > 0 for v in self.keypoints[::3]])
        return None

    def _compute_area(self):
        if self.bbox is not None:
            return self.bbox[-1] * self.bbox[-2]
        return None

    def to_dict(self):
        return self._to_dict_fields(
            ['image_id', 'id', 'category_id', 'bbox', 
             'area', 'segmentation', 'iscrowd', 
             'keypoints', 'num_keypoints'])

COCOBase

Bases: object

base class to facilitate conversion of COCO stuff to a dictionary.

TODO: refactor COCO classes to extend pycocotools' COCO

TODO: refactor COCO classes to use AttrDict

Source code in pycocowriter/coco.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class COCOBase(object):
    '''
    base class to facilitate conversion of COCO stuff to a dictionary.

    TODO: refactor COCO classes to extend pycocotools' COCO

    TODO: refactor COCO classes to use AttrDict
    '''
    def _to_dict_fields(self, fields:list[str]) -> dict:
        return {field: self.__dict__[field] for field in fields if self.__dict__[field] is not None}
    def to_dict(self):
        '''convert COCO object to dictionary'''
        raise NotImplementedError('must implement a to_dict method!')

to_dict()

convert COCO object to dictionary

Source code in pycocowriter/coco.py 
18
19
20
def to_dict(self):
    '''convert COCO object to dictionary'''
    raise NotImplementedError('must implement a to_dict method!')

COCOCategories

Bases: object

helper class to hold the index on categories so that we can find categories by name or by index

Parameters:

Name Type Description Default
categories list[COCOCategory] | None

existing list of COCOCategories if available

 None

Attributes:

Name Type Description
categories list[COCOCategory]

a list of unique categories
category_map dict[str, COCOCategory]

maps category names to categories
Source code in pycocowriter/coco.py 
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
235
236
237
238
239
240
241
242
243
244
245
246
class COCOCategories(object):
    '''
    helper class to hold the index on categories so that we can find categories by name or by index

    Parameters
    ----------
    categories: list[COCOCategory]
        existing list of COCOCategories if available

    Attributes
    ----------
    categories: list[COCOCategory]
        a list of unique categories
    category_map: dict[str, COCOCategory]
        maps category names to categories
    '''
    def __init__(self, categories: list[COCOCategory] | None = None):
        categories = categories or []
        # category ids MUST match their index+1 in the category list!
        self.categories = categories
        for i, category in enumerate(self.categories):
            assert category.id == i+1
        self.category_map = {category.name: category.id for category in self.categories}

    def add(self, label:str, keypoints:list[str]=None, skeleton:list[list[int]]=None) -> COCOCategory:
        '''
        Add a new category to the list.  Updates the map as well

        Parameters
        ----------
        label: str
            the string name of this category
        keypoints: list[str]
            the list of keypoint names for this category, if applicable
        skeleton: list[list[int]]
            the skeleton for this category, if applicable

        Returns
        -------
        category: COCOCategory
            returns the built COCOCategory
        '''
        if label not in self.category_map:
            category = COCOCategory(label, len(self.categories)+1, keypoints=keypoints, skeleton=skeleton)
            self.categories.append(category)
            self.category_map[self.categories[-1].name] = self.categories[-1].id
        return self.category_map[label]

    def __len__(self):
        return len(self.categories)

add(label, keypoints=None, skeleton=None)

Add a new category to the list. Updates the map as well

Parameters:

Name Type Description Default
label str

the string name of this category

 required
keypoints list[str]

the list of keypoint names for this category, if applicable

 None
skeleton list[list[int]]

the skeleton for this category, if applicable

 None

Returns:

Name Type Description
category COCOCategory

returns the built COCOCategory
Source code in pycocowriter/coco.py 
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
def add(self, label:str, keypoints:list[str]=None, skeleton:list[list[int]]=None) -> COCOCategory:
    '''
    Add a new category to the list.  Updates the map as well

    Parameters
    ----------
    label: str
        the string name of this category
    keypoints: list[str]
        the list of keypoint names for this category, if applicable
    skeleton: list[list[int]]
        the skeleton for this category, if applicable

    Returns
    -------
    category: COCOCategory
        returns the built COCOCategory
    '''
    if label not in self.category_map:
        category = COCOCategory(label, len(self.categories)+1, keypoints=keypoints, skeleton=skeleton)
        self.categories.append(category)
        self.category_map[self.categories[-1].name] = self.categories[-1].id
    return self.category_map[label]

COCOCategory

Bases: COCOBase

see https://cocodataset.org/#format-data

see https://github.com/facebookresearch/Detectron/issues/640

category {
    "id": int,
    "name": str,
    "supercategory": str,
    "keypoints": [str],
    "skeleton": [edge]
}

An edge is a tuple [a,b] where a,b are 1-indexed indices in the keypoints list. So if keypoints is ["a", "b"], then [1,2] is an edge between "a" and "b"

Source code in pycocowriter/coco.py 
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
class COCOCategory(COCOBase):
    '''
        see https://cocodataset.org/#format-data

        see https://github.com/facebookresearch/Detectron/issues/640

            category {
                "id": int,
                "name": str,
                "supercategory": str,
                "keypoints": [str],
                "skeleton": [edge]
            }

        An edge is a tuple [a,b] where a,b are 1-indexed indices in the keypoints list.  So if keypoints is ["a", "b"], then [1,2] is an edge between "a" and "b"
    '''
    def __init__(self, name:str, eye_d:int, supercategory:str=None, keypoints:list[str]=None, skeleton:list[list[int]]=None):
        self.name = name
        self.id = eye_d
        self.supercategory = supercategory
        self.keypoints = keypoints
        self.skeleton = skeleton

    def to_dict(self):
        return self._to_dict_fields(
            ['name', 'id', 'supercategory', 'keypoints', 'skeleton']
        )

COCOData

Bases: COCOBase

see https://cocodataset.org/#format-data

coco {
    "info": info,
    "images": [image],
    "annotations": [annotation],
    "licenses": [license],
    "categories": [category]
}
Source code in pycocowriter/coco.py 
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
class COCOData(COCOBase):
    '''
        see https://cocodataset.org/#format-data

            coco {
                "info": info,
                "images": [image],
                "annotations": [annotation],
                "licenses": [license],
                "categories": [category]
            }
    '''

    def __init__(self, info:COCOInfo, images:list[COCOImage], annotations:list[COCOAnnotation], licenses:list[COCOLicense], categories:list[COCOCategory]):
        self.info = info
        self.images = images
        self.annotations = annotations
        self.licenses = licenses
        self.categories = categories

    def to_dict(self) -> dict:
        return {
            'info': self.info.to_dict(),
            'images': [image.to_dict() for image in self.images],
            'annotations': [annotation.to_dict() for annotation in self.annotations],
            'licenses': [license.to_dict() for license in self.licenses],
            'categories': [category.to_dict() for category in self.categories]
        }

    def to_json(self, filename:str=None):
        '''
        dumps this COCO to json.  If a filename is provided, writes to disk, else, returns
        the string JSON

        Parameters
        ----------
        filename: str
            the optional location to which we should write the json

        Returns
        -------
        json: str | None
            if no filename is provided, returns the JSON COCO data as a string
        '''
        if filename is None:
            return json.dumps(self.to_dict(), cls=utils.NPEncoder)
        with open(filename, 'w') as f:
            json.dump(self.to_dict(), f, cls=utils.NPEncoder)

to_json(filename=None)

dumps this COCO to json. If a filename is provided, writes to disk, else, returns the string JSON

Parameters:

Name Type Description Default
filename str

the optional location to which we should write the json

 None

Returns:

Name Type Description
json str | None

if no filename is provided, returns the JSON COCO data as a string
Source code in pycocowriter/coco.py 
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
def to_json(self, filename:str=None):
    '''
    dumps this COCO to json.  If a filename is provided, writes to disk, else, returns
    the string JSON

    Parameters
    ----------
    filename: str
        the optional location to which we should write the json

    Returns
    -------
    json: str | None
        if no filename is provided, returns the JSON COCO data as a string
    '''
    if filename is None:
        return json.dumps(self.to_dict(), cls=utils.NPEncoder)
    with open(filename, 'w') as f:
        json.dump(self.to_dict(), f, cls=utils.NPEncoder)

COCOImage

Bases: COCOBase

see https://cocodataset.org/#format-data

image {
    "id": int,
    "width": int,
    "height": int,
    "file_name": str,
    "license": int,
    "flickr_url": str,
    "coco_url": str,
    "date_captured": datetime,
}
Source code in pycocowriter/coco.py 
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
class COCOImage(COCOBase):
    '''
        see https://cocodataset.org/#format-data

            image {
                "id": int,
                "width": int,
                "height": int,
                "file_name": str,
                "license": int,
                "flickr_url": str,
                "coco_url": str,
                "date_captured": datetime,
            }
    '''
    def __init__(self, 
                 eye_d:int, file_name:str, 
                 width:int=None, height:int=None, 
                 license:int=None, coco_url:str=None, 
                 date_captured:datetime.datetime=None,
                 discover_image_properties=True):
        self.id = eye_d
        self.file_name = file_name
        self.width = width
        self.height = height
        if ((self.width is None) or (self.height is None)) and discover_image_properties:
            self.compute_width_and_height()
        self.license = license
        self.coco_url = coco_url
        self.date_captured = date_captured

    def compute_width_and_height(self):
        with Image.open(self.file_name) as im:
            self.width, self.height = im.size

    def to_dict(self) -> dict:
        the_dict = self._to_dict_fields(
            ['id', 'file_name', 'width', 'height', 'license', 'coco_url']
        )
        if self.date_captured:
            the_dict['date_captured'] = self.date_captured.isoformat()
        return the_dict

COCOImages

Bases: object

helper class to hold the index on images so that we can find images by name or by index

Parameters:

Name Type Description Default
images list[COCOImage] | None

existing list of COCOImages if available

 None

Attributes:

Name Type Description
images list[COCOImage]

a list of unique images
image_map dict[str, COCOImage]

maps image names to images
Source code in pycocowriter/coco.py 
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
class COCOImages(object):
    '''
    helper class to hold the index on images so that we can find images by name or by index

    Parameters
    ----------
    images: list[COCOImage]
        existing list of COCOImages if available

    Attributes
    ----------
    images: list[COCOImage]
        a list of unique images
    image_map: dict[str, COCOImage]
        maps image names to images
    '''
    def __init__(self, images: list[COCOImage] | None = None):
        images = images or []
        # image ids MUST match their index+1 in the image list!
        self.images = images
        for i, image in enumerate(self.images):
            assert image.id == i+1
        self.image_map = {image.filename: image.id for image in self.images}

    def add(self, filename:str, width:int=None, height:int=None, 
            url:str=None, license:int=None, date_captured:datetime.datetime=None):
        '''
        Add a new image to the list.  Updates the map as well

        Parameters
        ----------
        filename: str
            the filename of this image
        width: int
            the width of this image, if known
        height: int
            the height of this image, if known
        url: str
            the url at which this image can be downloaded
        license: int
            the index in the list of COCOLicense with the applicable license
        date_captured: datetime.datetime
            the date and time when the image was captured

        Returns
        -------
        image: COCOImage
            returns the built COCOImage
        '''
        if filename not in self.image_map:
            image = COCOImage(len(self.images)+1, filename, 
                              width=width, height=height, coco_url=url, 
                              license=license, date_captured=date_captured)
            self.images.append(image)
            self.image_map[self.images[-1].file_name] = self.images[-1].id
        return self.image_map[filename]

    def __len__(self):
        return len(self.images)

add(filename, width=None, height=None, url=None, license=None, date_captured=None)

Add a new image to the list. Updates the map as well

Parameters:

Name Type Description Default
filename str

the filename of this image

 required
width int

the width of this image, if known

 None
height int

the height of this image, if known

 None
url str

the url at which this image can be downloaded

 None
license int

the index in the list of COCOLicense with the applicable license

 None
date_captured datetime

the date and time when the image was captured

 None

Returns:

Name Type Description
image COCOImage

returns the built COCOImage
Source code in pycocowriter/coco.py 
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
def add(self, filename:str, width:int=None, height:int=None, 
        url:str=None, license:int=None, date_captured:datetime.datetime=None):
    '''
    Add a new image to the list.  Updates the map as well

    Parameters
    ----------
    filename: str
        the filename of this image
    width: int
        the width of this image, if known
    height: int
        the height of this image, if known
    url: str
        the url at which this image can be downloaded
    license: int
        the index in the list of COCOLicense with the applicable license
    date_captured: datetime.datetime
        the date and time when the image was captured

    Returns
    -------
    image: COCOImage
        returns the built COCOImage
    '''
    if filename not in self.image_map:
        image = COCOImage(len(self.images)+1, filename, 
                          width=width, height=height, coco_url=url, 
                          license=license, date_captured=date_captured)
        self.images.append(image)
        self.image_map[self.images[-1].file_name] = self.images[-1].id
    return self.image_map[filename]

COCOInfo

Bases: COCOBase

see https://cocodataset.org/#format-data

info {
    "year": int,
    "version": str,
    "description": str,
    "contributor": str,
    "url": str,
    "date_created": datetime,
}
Source code in pycocowriter/coco.py 
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
class COCOInfo(COCOBase):
    '''
        see https://cocodataset.org/#format-data

            info {
                "year": int,
                "version": str,
                "description": str,
                "contributor": str,
                "url": str,
                "date_created": datetime,
            }
    '''
    def __init__(self, year:int=None, version:str=None, description:str=None, contributor:str=None, url:str=None, date_created:datetime.datetime=None):
        self.year = year
        self.version = version
        self.description = description
        self.contributor = contributor
        self.url = url
        self.date_created = date_created

    def to_dict(self) -> dict:
        the_dict = self._to_dict_fields(
            ['year', 'version', 'description', 'contributor', 'url']
        )
        if self.date_created:
            the_dict['date_created'] = self.date_created.isoformat()
        return the_dict

COCOLicense

Bases: COCOBase

see https://cocodataset.org/#format-data

license {
    "id": int,
    "name": str,
    "url": str,
}
Source code in pycocowriter/coco.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
class COCOLicense(COCOBase):
    '''
        see https://cocodataset.org/#format-data

            license {
                "id": int,
                "name": str,
                "url": str,
            }
    '''
    def __init__(self, name:str, eye_d:int, url:str=None):
        self.name = name
        self.id = eye_d
        self.url = url

    def to_dict(self):
        return self._to_dict_fields(
            ['name', 'id', 'url']
        )