def test_bbox_mask():
# default config for random bbox mask
cfg = dict(img_shape=(256, 256),
max_bbox_shape=100,
max_bbox_delta=10,
min_margin=10)
bbox = random_bbox(**cfg)
mask_bbox = bbox2mask(cfg['img_shape'], bbox)
assert mask_bbox.shape == (256, 256, 1)
zero_area = np.sum((mask_bbox == 0).astype(np.uint8))
ones_area = np.sum((mask_bbox == 1).astype(np.uint8))
assert zero_area + ones_area == 256 * 256
assert mask_bbox.dtype == np.uint8
with pytest.raises(ValueError):
cfg_ = cfg.copy()
cfg_['max_bbox_shape'] = 300
bbox = random_bbox(**cfg_)
with pytest.raises(ValueError):
cfg_ = cfg.copy()
cfg_['max_bbox_delta'] = 300
bbox = random_bbox(**cfg_)
with pytest.raises(ValueError):
cfg_ = cfg.copy()
cfg_['max_bbox_shape'] = 254
bbox = random_bbox(**cfg_)
cfg_ = cfg.copy()
cfg_['max_bbox_delta'] = 1
bbox = random_bbox(**cfg_)
mask_bbox = bbox2mask(cfg['img_shape'], bbox)
assert mask_bbox.shape == (256, 256, 1)
def __call__(self, results):
"""Call function.
Args:
results (dict): A dict containing the necessary information and
data for augmentation.
Returns:
dict: A dict containing the processed data and information.
"""
if self.mask_mode == 'bbox':
mask_bbox = random_bbox(**self.mask_config)
mask = bbox2mask(self.mask_config['img_shape'], mask_bbox)
results['mask_bbox'] = mask_bbox
elif self.mask_mode == 'irregular':
mask = get_irregular_mask(**self.mask_config)
elif self.mask_mode == 'set':
mask = self._get_random_mask_from_set()
elif self.mask_mode == 'ff':
mask = brush_stroke_mask(**self.mask_config)
elif self.mask_mode == 'file':
mask = self._get_mask_from_file(results['mask_path'])
else:
raise NotImplementedError(
f'Mask mode {self.mask_mode} has not been implemented.')
results['mask'] = mask
return results