-
Notifications
You must be signed in to change notification settings - Fork 3
/
vsrl_utils.py
155 lines (138 loc) · 5.9 KB
/
vsrl_utils.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
# AUTORIGHTS
# ---------------------------------------------------------
# Copyright (c) 2016, Saurabh Gupta
#
# This file is part of the VCOCO dataset hooks and is available
# under the terms of the Simplified BSD License provided in
# LICENSE. Please retain this notice and LICENSE if you use
# this file (or any portion of it) in your project.
# ---------------------------------------------------------
# vsrl_data is a dictionary for each action class:
# image_id - Nx1
# ann_id - Nx1
# label - Nx1
# action_name - string
# role_name - ['agent', 'obj', 'instr']
# role_object_id - N x K matrix, obviously [:,0] is same as ann_id
import numpy as np
from pycocotools.coco import COCO
import os, json
coco = []
def get_data_dir():
this_dir = os.path.dirname(__file__)
dir_name = os.path.join(this_dir, 'data')
return dir_name
def load_coco(dir_name=None):
global coco
if dir_name is None:
dir_name = get_data_dir()
if coco == []:
coco = COCO(os.path.join(dir_name, 'instances_vcoco_all_2014.json'))
return coco
def load_vcoco(imset, dir_name=None):
if dir_name is None:
dir_name = get_data_dir()
print(os.path.join(dir_name, 'vcoco', imset + '.json'))
with open(os.path.join(dir_name, 'vcoco', imset + '.json'), 'rt') as f:
vsrl_data = json.load(f)
vsrl_data = unicode_to_str(vsrl_data)
for i in range(len(vsrl_data)):
vsrl_data[i]['role_object_id'] = \
np.array(vsrl_data[i]['role_object_id']).reshape((len(vsrl_data[i]['role_name']),-1)).T
for j in ['ann_id', 'label', 'image_id']:
vsrl_data[i][j] = np.array(vsrl_data[i][j]).reshape((-1,1))
return vsrl_data
def coco_ann_id_to_image_id(ann_id, coco):
ann = coco.loadAnns(ann_id);
ann_id_r = [a['id'] for a in ann]
image_id_r = [a['image_id'] for a in ann]
assert ann_id_r == ann_id, \
'oops coco returned different ann_id''s in different order'
return image_id_r
# Attach relevant boxes
def all_relevant_objects(vsrl_data_all, coco):
vsrl_relevant_objects = []
for i in xrange(len(vsrl_data_all)):
v_i = vsrl_data_all[i]
bbox = []; image_id = [];
num_roles = len(v_i['role_name'])
for j in range(num_roles):
# print v_i['action_name'], v_i['include'][j];
bbox_j = np.NaN*np.ones((0,4*num_roles));
image_id_j = np.NaN*np.ones((0,1))
if v_i['include'][j] != []:
cat_ids = coco.getCatIds(catNms=v_i['include'][j])
ann_list = coco.getAnnIds(imgIds=np.unique(v_i['image_id']*1).tolist(),
iscrowd=False, catIds=cat_ids)
anns = coco.loadAnns(ann_list)
bbox_j = np.NaN*np.zeros((len(anns), 4*num_roles))
bbox_j[:,4*j:4*j+4] = xyhw_to_xyxy(np.vstack([np.array(a['bbox']).reshape((1,4)) for a in anns]))
image_id_j = np.array(coco_ann_id_to_image_id(ann_list, coco)).reshape(-1,1)
bbox.append(bbox_j)
image_id.append(image_id_j)
image_id = np.concatenate(image_id, axis=0)
bbox = np.concatenate(bbox, axis=0)
vsrl_relevant_objects.append({'image_id': image_id, 'bbox': bbox})
return vsrl_relevant_objects
def attach_unlabelled(vsrl_data, coco):
"""
def vsrl_data = attach_unlabelled(vsrl_data, coco)
"""
anns = coco.loadAnns(
coco.getAnnIds(imgIds=np.unique(vsrl_data['image_id']).tolist(),
iscrowd=False, catIds=1));
ann_id = [a['id'] for a in anns]
hard_ann_id = list(set(ann_id) - set(vsrl_data['ann_id'].ravel().tolist()))
hard_image_id = coco_ann_id_to_image_id(hard_ann_id, coco);
vsrl_data['image_id'] = np.vstack((vsrl_data['image_id'],
np.array(hard_image_id).reshape((-1,1))))
vsrl_data['ann_id'] = np.vstack((vsrl_data['ann_id'],
np.array(hard_ann_id).reshape((-1,1))))
vsrl_data['role_object_id'] = np.vstack((vsrl_data['role_object_id'],
np.zeros((len(hard_image_id), vsrl_data['role_object_id'].shape[1]))))
vsrl_data['role_object_id'][vsrl_data['label'].shape[0]:,0] = np.array(hard_ann_id).reshape((-1))
vsrl_data['label'] = np.vstack((vsrl_data['label'],
-1*np.ones((len(hard_image_id), vsrl_data['label'].shape[1]))))
return vsrl_data
def remove_negative(vsrl_data):
"""
def vsrl_data = remove_negative(vsrl_data)
Remove things that are labelled as a negative
"""
to_keep = vsrl_data['label'] != 0
for i in vsrl_data.keys():
if type(vsrl_data[i]) == np.ndarray:
vsrl_data[i] = vsrl_data[i][to_keep, :]
return vsrl_data
def xyhw_to_xyxy(bbox):
out = bbox.copy()
out[:, [2, 3]] = bbox[:, [0,1]] + bbox[:, [2,3]];
return out
def attach_gt_boxes(vsrl_data, coco):
"""
def vsrl_data = attach_gt_boxes(vsrl_data, coco)
"""
anns = coco.loadAnns(vsrl_data['ann_id'].ravel().tolist());
bbox = np.vstack([np.array(a['bbox']).reshape((1,4)) for a in anns])
vsrl_data['bbox'] = xyhw_to_xyxy(bbox)
vsrl_data['role_bbox'] = \
np.nan*np.zeros((vsrl_data['role_object_id'].shape[0], \
4*vsrl_data['role_object_id'].shape[1]), dtype=np.float)
# Get boxes for the role objects
for i in range(vsrl_data['role_object_id'].shape[1]):
has_role = np.where(vsrl_data['role_object_id'][:,i] > 0)[0]
if has_role.size > 0:
anns = coco.loadAnns(vsrl_data['role_object_id'][has_role, i].ravel().tolist());
bbox = np.vstack([np.array(a['bbox']).reshape((1,4)) for a in anns])
bbox = xyhw_to_xyxy(bbox)
vsrl_data['role_bbox'][has_role, 4*i:4*(i+1)] = bbox;
return vsrl_data
def unicode_to_str(input):
if isinstance(input, dict):
return {unicode_to_str(key):unicode_to_str(value) for key,value in input.iteritems()}
elif isinstance(input, list):
return [unicode_to_str(element) for element in input]
elif isinstance(input, unicode):
return input.encode('utf-8')
else:
return input