-
Notifications
You must be signed in to change notification settings - Fork 605
/
coco.py
156 lines (118 loc) · 5.63 KB
/
coco.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
import os
from .base_image_dataset import BaseImageDataset
from ltr.data.image_loader import jpeg4py_loader
import torch
from pycocotools.coco import COCO
import random
from collections import OrderedDict
from ltr.admin.environment import env_settings
class MSCOCO(BaseImageDataset):
""" The COCO object detection dataset.
Publication:
Microsoft COCO: Common Objects in Context.
Tsung-Yi Lin, Michael Maire, Serge J. Belongie, Lubomir D. Bourdev, Ross B. Girshick, James Hays, Pietro Perona,
Deva Ramanan, Piotr Dollar and C. Lawrence Zitnick
ECCV, 2014
https://arxiv.org/pdf/1405.0312.pdf
Download the images along with annotations from http://cocodataset.org/#download. The root folder should be
organized as follows.
- coco_root
- annotations
- instances_train2014.json
- instances_train2017.json
- images
- train2014
- train2017
Note: You also have to install the coco pythonAPI from https://github.com/cocodataset/cocoapi.
"""
def __init__(self, root=None, image_loader=jpeg4py_loader, data_fraction=None, min_area=None,
split="train", version="2014"):
"""
args:
root - path to coco root folder
image_loader (jpeg4py_loader) - The function to read the images. jpeg4py (https://github.com/ajkxyz/jpeg4py)
is used by default.
data_fraction - Fraction of dataset to be used. The complete dataset is used by default
min_area - Objects with area less than min_area are filtered out. Default is 0.0
split - 'train' or 'val'.
version - version of coco dataset (2014 or 2017)
"""
root = env_settings().coco_dir if root is None else root
super().__init__('COCO', root, image_loader)
self.img_pth = os.path.join(root, 'images/{}{}/'.format(split, version))
self.anno_path = os.path.join(root, 'annotations/instances_{}{}.json'.format(split, version))
self.coco_set = COCO(self.anno_path)
self.cats = self.coco_set.cats
self.class_list = self.get_class_list() # the parent class thing would happen in the sampler
self.image_list = self._get_image_list(min_area=min_area)
if data_fraction is not None:
self.image_list = random.sample(self.image_list, int(len(self.image_list) * data_fraction))
self.im_per_class = self._build_im_per_class()
def _get_image_list(self, min_area=None):
ann_list = list(self.coco_set.anns.keys())
image_list = [a for a in ann_list if self.coco_set.anns[a]['iscrowd'] == 0]
if min_area is not None:
image_list = [a for a in image_list if self.coco_set.anns[a]['area'] > min_area]
return image_list
def get_num_classes(self):
return len(self.class_list)
def get_name(self):
return 'coco'
def has_class_info(self):
return True
def has_segmentation_info(self):
return True
def get_class_list(self):
class_list = []
for cat_id in self.cats.keys():
class_list.append(self.cats[cat_id]['name'])
return class_list
def _build_im_per_class(self):
im_per_class = {}
for i, im in enumerate(self.image_list):
class_name = self.cats[self.coco_set.anns[im]['category_id']]['name']
if class_name not in im_per_class:
im_per_class[class_name] = [i]
else:
im_per_class[class_name].append(i)
return im_per_class
def get_images_in_class(self, class_name):
return self.im_per_class[class_name]
def get_image_info(self, im_id):
anno = self._get_anno(im_id)
bbox = torch.Tensor(anno['bbox']).view(4,)
mask = torch.Tensor(self.coco_set.annToMask(anno))
valid = (bbox[2] > 0) & (bbox[3] > 0)
visible = valid.clone().byte()
return {'bbox': bbox, 'mask': mask, 'valid': valid, 'visible': visible}
def _get_anno(self, im_id):
anno = self.coco_set.anns[self.image_list[im_id]]
return anno
def _get_image(self, im_id):
path = self.coco_set.loadImgs([self.coco_set.anns[self.image_list[im_id]]['image_id']])[0]['file_name']
img = self.image_loader(os.path.join(self.img_pth, path))
return img
def get_meta_info(self, im_id):
try:
cat_dict_current = self.cats[self.coco_set.anns[self.image_list[im_id]]['category_id']]
object_meta = OrderedDict({'object_class_name': cat_dict_current['name'],
'motion_class': None,
'major_class': cat_dict_current['supercategory'],
'root_class': None,
'motion_adverb': None})
except:
object_meta = OrderedDict({'object_class_name': None,
'motion_class': None,
'major_class': None,
'root_class': None,
'motion_adverb': None})
return object_meta
def get_class_name(self, im_id):
cat_dict_current = self.cats[self.coco_set.anns[self.image_list[im_id]]['category_id']]
return cat_dict_current['name']
def get_image(self, image_id, anno=None):
frame = self._get_image(image_id)
if anno is None:
anno = self.get_image_info(image_id)
object_meta = self.get_meta_info(image_id)
return frame, anno, object_meta