Skip to content
Snippets Groups Projects
Unverified Commit 3d790bcb authored by Kai Chen's avatar Kai Chen Committed by GitHub
Browse files

Merge pull request #28 from myownskyW7/dev

update coco dataset to support proposals with shapes (n, 5)
parents f29f020c ad5c27e1
No related branches found
No related tags found
No related merge requests found
...@@ -203,13 +203,22 @@ class CocoDataset(Dataset): ...@@ -203,13 +203,22 @@ class CocoDataset(Dataset):
# load proposals if necessary # load proposals if necessary
if self.proposals is not None: if self.proposals is not None:
proposals = self.proposals[idx][:self.num_max_proposals, :4] proposals = self.proposals[idx][:self.num_max_proposals]
# TODO: Handle empty proposals properly. Currently images with # TODO: Handle empty proposals properly. Currently images with
# no proposals are just ignored, but they can be used for # no proposals are just ignored, but they can be used for
# training in concept. # training in concept.
if len(proposals) == 0: if len(proposals) == 0:
idx = self._rand_another(idx) idx = self._rand_another(idx)
continue continue
if not (proposals.shape[1] == 4 or proposals.shape[1] == 5):
raise AssertionError(
'proposals should have shapes (n, 4) or (n, 5), '
'but found {}'.format(proposals.shape))
if proposals.shape[1] == 5:
scores = proposals[:, 4]
proposals = proposals[:, :4]
else:
scores = None
ann = self._parse_ann_info(ann_info, self.with_mask) ann = self._parse_ann_info(ann_info, self.with_mask)
gt_bboxes = ann['bboxes'] gt_bboxes = ann['bboxes']
...@@ -228,6 +237,8 @@ class CocoDataset(Dataset): ...@@ -228,6 +237,8 @@ class CocoDataset(Dataset):
if self.proposals is not None: if self.proposals is not None:
proposals = self.bbox_transform(proposals, img_shape, proposals = self.bbox_transform(proposals, img_shape,
scale_factor, flip) scale_factor, flip)
proposals = np.hstack([proposals, scores[:, None]
]) if scores is not None else proposals
gt_bboxes = self.bbox_transform(gt_bboxes, img_shape, scale_factor, gt_bboxes = self.bbox_transform(gt_bboxes, img_shape, scale_factor,
flip) flip)
gt_bboxes_ignore = self.bbox_transform(gt_bboxes_ignore, img_shape, gt_bboxes_ignore = self.bbox_transform(gt_bboxes_ignore, img_shape,
...@@ -263,8 +274,12 @@ class CocoDataset(Dataset): ...@@ -263,8 +274,12 @@ class CocoDataset(Dataset):
"""Prepare an image for testing (multi-scale and flipping)""" """Prepare an image for testing (multi-scale and flipping)"""
img_info = self.img_infos[idx] img_info = self.img_infos[idx]
img = mmcv.imread(osp.join(self.img_prefix, img_info['file_name'])) img = mmcv.imread(osp.join(self.img_prefix, img_info['file_name']))
proposal = (self.proposals[idx][:, :4] if self.proposals is not None:
if self.proposals is not None else None) proposal = self.proposals[idx][:self.num_max_proposals]
if not (proposal.shape[1] == 4 or proposal.shape[1] == 5):
raise AssertionError(
'proposals should have shapes (n, 4) or (n, 5), '
'but found {}'.format(proposal.shape))
def prepare_single(img, scale, flip, proposal=None): def prepare_single(img, scale, flip, proposal=None):
_img, img_shape, pad_shape, scale_factor = self.img_transform( _img, img_shape, pad_shape, scale_factor = self.img_transform(
...@@ -277,8 +292,15 @@ class CocoDataset(Dataset): ...@@ -277,8 +292,15 @@ class CocoDataset(Dataset):
scale_factor=scale_factor, scale_factor=scale_factor,
flip=flip) flip=flip)
if proposal is not None: if proposal is not None:
if proposal.shape[1] == 5:
score = proposal[:, 4]
proposal = proposal[:, :4]
else:
score = None
_proposal = self.bbox_transform(proposal, img_shape, _proposal = self.bbox_transform(proposal, img_shape,
scale_factor, flip) scale_factor, flip)
_proposal = np.hstack([_proposal, score[:, None]
]) if score is not None else _proposal
_proposal = to_tensor(_proposal) _proposal = to_tensor(_proposal)
else: else:
_proposal = None _proposal = None
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment