def collate_fn_testing(data):
im_batch, im_scale_batch, anchors, _, _, _ = \
data_layer_keep_aspect_ratio_batch(data, is_training=False)
input = torch.stack(everything2tensor(im_batch))
gt_boxes_list = [d[2]['gt_boxes'] for d in data]
im_scale_list = [d[1] for d in data]
image_ids = [d[3] for d in data]
return input, anchors, im_scale_list, image_ids, gt_boxes_list
def collate_fn(data):
# print ('collate_fn:', len(data))
im_batch, im_scale_batch, anchors, rpn_targets, inst_masks_batch, mask_batch = \
data_layer_keep_aspect_ratio_batch(data, is_training=True)
input = torch.stack(everything2tensor(im_batch))
# input = np.stack(im_batch)
gt_boxes_list = [d[2]['gt_boxes'] for d in data]
image_ids = [d[3] for d in data]
im_scale_list = [d[1] for d in data]
labels = torch.stack(
everything2tensor([d[2]['rpn_targets'][0] for d in data]))
# label_weights = torch.stack(everything2tensor(rpn_targets['labels_weights_batch']))
bbox_targets = torch.stack(
everything2tensor([d[2]['rpn_targets'][1] for d in data]))
bbox_inside_weights = torch.stack(
everything2tensor([d[2]['rpn_targets'][2] for d in data]))
rpn_targets = [labels, bbox_targets, bbox_inside_weights]
mask = torch.stack(everything2tensor(mask_batch))
inst_masks = everything2tensor(
inst_masks_batch) # a list of (N, H, W). N can be changable
return input, anchors, im_scale_list, image_ids, gt_boxes_list, rpn_targets, inst_masks, mask
def __getitem__(self, i):
img_id = self._imgs[i][0]
img_name = self._imgs[i][1]['file_name']
split = img_name.split('_')[1]
img_name = os.path.join(self._data_dir, split, img_name)
coco = self._cocos[self._imgs[i][2]]
# if self._split == 'trainval2014' and not self._imgs[i][2]:
# coco = self._cocos[1]
height, width = self._imgs[i][1]['height'], self._imgs[i][1]['width']
bboxes, classes, inst_masks, mask = self._get_coco_masks(
coco, img_id, height, width, img_name)
# img = np.array(Image.open(img_name))
# do some preprocessings here
im, TARGETS, inst_masks, mask, ori_im, ANNOTATIONS = \
self._data_handler(img_name, bboxes, classes, inst_masks, mask, self._is_training, self.ANCHORS)
im = np.transpose(im, [2, 0, 1]) # c, h, w
if self._is_training:
ih, iw = im.shape[1], im.shape[2]
# masks = np.transpose(masks, [0, 3, 1, 2]) # n,
# mask = mask[np.newaxis, :, :] # 1xhxw
downsampled_mask = []
for _, stride in enumerate(cfg.strides):
assert ih % stride == iw % stride == 0, '{} {} {}'.format(
ih, iw, stride)
h, w = ih / stride, iw / stride
downsampled_mask.append(
cv2.resize(mask, (h, w),
interpolation=cv2.INTER_NEAREST).astype(
np.int64))
# to tensor
im = everything2tensor(im.copy())
TARGETS = everything2tensor(TARGETS)
inst_masks = everything2tensor(inst_masks.astype(np.int64).copy())
mask = everything2tensor(mask.astype(np.int64).copy())
downsampled_mask = everything2tensor(downsampled_mask)
else:
"""testing"""
im = everything2tensor(im.copy())
TARGETS = inst_masks = mask = downsampled_mask = ANNOTATIONS = []
# return image id
img_id = self._imgs[i][0]
return im, TARGETS, inst_masks, mask, downsampled_mask, ori_im, ANNOTATIONS, img_id
def __getitem__(self, i):
img_name = self._image_list[i]
im = cv2.imread(img_name)
height, width, _ = im.shape
# empty annotations
bboxes = np.asarray([[0, 0, width - 1, height - 1]], dtype=np.float32)
classes = np.asarray((-1, ), dtype=np.int32)
inst_masks = np.zeros([1, height, width], dtype=np.int32)
mask = np.zeros([height, width], dtype=np.int32)
im, TARGETS, inst_masks, mask, ori_im, ANNOTATIONS = \
self._data_handler(img_name, bboxes, classes, inst_masks, mask, self._is_training, self.ANCHORS)
im = np.transpose(im, [2, 0, 1]) # c, h, w
im = everything2tensor(im.copy())
TARGETS = inst_masks = mask = downsampled_mask = []
img_id = os.path.split(img_name)[1]
return im, TARGETS, inst_masks, mask, downsampled_mask, ori_im, ANNOTATIONS, img_id