in vilbert/datasets/_image_features_reader.py [0:0]
def __getitem__(self, image_id):
image_id = str(image_id).encode()
index = self._image_ids.index(image_id)
if self._in_memory:
# Load features during first epoch, all not loaded together as it
# has a slow start.
if self.features[index] is not None:
features = self.features[index]
num_boxes = self.num_boxes[index]
image_location = self.boxes[index]
image_location_ori = self.boxes_ori[index]
else:
with self.env.begin(write=False) as txn:
item = pickle.loads(txn.get(image_id))
image_id = item["image_id"]
image_h = int(item["image_h"])
image_w = int(item["image_w"])
# num_boxes = int(item['num_boxes'])
# features = np.frombuffer(base64.b64decode(item["features"]), dtype=np.float32).reshape(num_boxes, 2048)
# boxes = np.frombuffer(base64.b64decode(item['boxes']), dtype=np.float32).reshape(num_boxes, 4)
features = item["features"].reshape(-1, 2048)
boxes = item["boxes"].reshape(-1, 4)
num_boxes = features.shape[0]
g_feat = np.sum(features, axis=0) / num_boxes
num_boxes = num_boxes + 1
features = np.concatenate(
[np.expand_dims(g_feat, axis=0), features], axis=0
)
self.features[index] = features
image_location = np.zeros((boxes.shape[0], 5), dtype=np.float32)
image_location[:, :4] = boxes
image_location[:, 4] = (
(image_location[:, 3] - image_location[:, 1])
* (image_location[:, 2] - image_location[:, 0])
/ (float(image_w) * float(image_h))
)
image_location_ori = copy.deepcopy(image_location)
image_location[:, 0] = image_location[:, 0] / float(image_w)
image_location[:, 1] = image_location[:, 1] / float(image_h)
image_location[:, 2] = image_location[:, 2] / float(image_w)
image_location[:, 3] = image_location[:, 3] / float(image_h)
g_location = np.array([0, 0, 1, 1, 1])
image_location = np.concatenate(
[np.expand_dims(g_location, axis=0), image_location], axis=0
)
self.boxes[index] = image_location
g_location_ori = np.array(
[0, 0, image_w, image_h, image_w * image_h]
)
image_location_ori = np.concatenate(
[np.expand_dims(g_location_ori, axis=0), image_location_ori],
axis=0,
)
self.boxes_ori[index] = image_location_ori
self.num_boxes[index] = num_boxes
else:
# Read chunk from file everytime if not loaded in memory.
with self.env.begin(write=False) as txn:
item = pickle.loads(txn.get(image_id))
image_id = item["image_id"]
image_h = int(item["image_h"])
image_w = int(item["image_w"])
# num_boxes = int(item['num_boxes'])
# features = np.frombuffer(base64.b64decode(item["features"]), dtype=np.float32).reshape(num_boxes, 2048)
# boxes = np.frombuffer(base64.b64decode(item['boxes']), dtype=np.float32).reshape(num_boxes, 4)
features = item["features"].reshape(-1, 2048)
boxes = item["boxes"].reshape(-1, 4)
num_boxes = features.shape[0]
g_feat = np.sum(features, axis=0) / num_boxes
num_boxes = num_boxes + 1
features = np.concatenate(
[np.expand_dims(g_feat, axis=0), features], axis=0
)
image_location = np.zeros((boxes.shape[0], 5), dtype=np.float32)
image_location[:, :4] = boxes
image_location[:, 4] = (
(image_location[:, 3] - image_location[:, 1])
* (image_location[:, 2] - image_location[:, 0])
/ (float(image_w) * float(image_h))
)
image_location_ori = copy.deepcopy(image_location)
image_location[:, 0] = image_location[:, 0] / float(image_w)
image_location[:, 1] = image_location[:, 1] / float(image_h)
image_location[:, 2] = image_location[:, 2] / float(image_w)
image_location[:, 3] = image_location[:, 3] / float(image_h)
g_location = np.array([0, 0, 1, 1, 1])
image_location = np.concatenate(
[np.expand_dims(g_location, axis=0), image_location], axis=0
)
g_location_ori = np.array([0, 0, image_w, image_h, image_w * image_h])
image_location_ori = np.concatenate(
[np.expand_dims(g_location_ori, axis=0), image_location_ori], axis=0
)
return features, num_boxes, image_location, image_location_ori