def _load_masks(self, results):
ann_info = results['ann_info']
h, w = results['img_info']['height'], results['img_info']['width']
gt_masks = ann_info['masks']
if self.poly2mask:
gt_masks = BitmapMasks(
[self._poly2mask(mask, h, w) for mask in gt_masks], h, w)
else:
gt_masks = PolygonMasks(
[self.process_polygons(polygons) for polygons in gt_masks], h,
w)
gt_masks_ignore = ann_info.get('masks_ignore', None)
if gt_masks_ignore is not None:
if self.poly2mask:
gt_masks_ignore = BitmapMasks(
[self._poly2mask(mask, h, w) for mask in gt_masks_ignore],
h, w)
else:
gt_masks_ignore = PolygonMasks([
self.process_polygons(polygons)
for polygons in gt_masks_ignore
], h, w)
results['gt_masks_ignore'] = gt_masks_ignore
results['mask_fields'].append('gt_masks_ignore')
results['gt_masks'] = gt_masks
results['mask_fields'].append('gt_masks')
return results
def test_polygon_mask_init():
# init with empty masks
raw_masks = []
polygon_masks = BitmapMasks(raw_masks, 28, 28)
assert len(polygon_masks) == 0
assert polygon_masks.height == 28
assert polygon_masks.width == 28
# init with masks contain 3 instances
raw_masks = dummy_raw_polygon_masks((3, 28, 28))
polygon_masks = PolygonMasks(raw_masks, 28, 28)
assert isinstance(polygon_masks.masks, list)
assert isinstance(polygon_masks.masks[0], list)
assert isinstance(polygon_masks.masks[0][0], np.ndarray)
assert len(polygon_masks) == 3
assert polygon_masks.height == 28
assert polygon_masks.width == 28
assert polygon_masks.to_ndarray().shape == (3, 28, 28)
# init with raw masks of unsupported type
with pytest.raises(AssertionError):
raw_masks = [[[]]]
PolygonMasks(raw_masks, 28, 28)
raw_masks = [dummy_raw_polygon_masks((3, 28, 28))]
PolygonMasks(raw_masks, 28, 28)
def _load_masks(self, results):
"""Private function to load mask annotations.
Args:
results (dict): Result dict from :obj:`mmdet.CustomDataset`.
Returns:
dict: The dict contains loaded mask annotations.
If ``self.poly2mask`` is set ``True``, `gt_mask` will contain
:obj:`PolygonMasks`. Otherwise, :obj:`BitmapMasks` is used.
"""
h, w = results['img_info']['height'], results['img_info']['width']
gt_masks = results['ann_info']['masks']
if self.poly2mask:
gt_masks = BitmapMasks(
[self._poly2mask(mask, h, w) for mask in gt_masks], h, w)
elif self.OC2mask:
gt_masks = BitmapMasks(
[self.ochuman2mask(mask) for mask in gt_masks], h, w)
else:
gt_masks = PolygonMasks(
[self.process_polygons(polygons) for polygons in gt_masks], h,
w)
results['gt_masks'] = gt_masks
results['mask_fields'].append('gt_masks')
return results
def test_bitmap_mask_init():
# init with empty ndarray masks
raw_masks = np.empty((0, 28, 28), dtype=np.uint8)
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
assert len(bitmap_masks) == 0
assert bitmap_masks.height == 28
assert bitmap_masks.width == 28
# init with empty list masks
raw_masks = []
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
assert len(bitmap_masks) == 0
assert bitmap_masks.height == 28
assert bitmap_masks.width == 28
# init with ndarray masks contain 3 instances
raw_masks = dummy_raw_bitmap_masks((3, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
assert len(bitmap_masks) == 3
assert bitmap_masks.height == 28
assert bitmap_masks.width == 28
# init with list masks contain 3 instances
raw_masks = [dummy_raw_bitmap_masks((28, 28)) for _ in range(3)]
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
assert len(bitmap_masks) == 3
assert bitmap_masks.height == 28
assert bitmap_masks.width == 28
# init with raw masks of unsupported type
with pytest.raises(AssertionError):
raw_masks = [[dummy_raw_bitmap_masks((28, 28))]]
BitmapMasks(raw_masks, 28, 28)
def test_bitmap_mask_get_bboxes():
# resize with empty bitmap masks
raw_masks = dummy_raw_bitmap_masks((0, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
bboxes = bitmap_masks.get_bboxes()
assert len(bboxes) == 0
# resize with bitmap masks contain 1 instances
raw_masks = np.array([[[0, 0, 0, 0, 0, 0, 0, 0], [0, 1, 1, 1, 0, 0, 0, 0],
[0, 0, 1, 1, 0, 0, 0, 0], [0, 0, 1, 1, 1, 0, 0, 0],
[0, 0, 1, 1, 1, 1, 0, 0], [0, 0, 1, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0,
0]]])
bitmap_masks = BitmapMasks(raw_masks, 8, 8)
bboxes = bitmap_masks.get_bboxes()
assert len(bboxes) == 1
truth = np.array([[1, 1, 6, 6]])
assert (bboxes == truth).all()
# resize to non-square
raw_masks = np.array([[[1, 1, 0, 0, 0, 0, 0, 0], [0, 0, 1, 1, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 1, 0, 0], [0, 0, 0, 0, 0, 0, 0,
0]]])
bitmap_masks = BitmapMasks(raw_masks, 4, 8)
bboxes = bitmap_masks.get_bboxes()
truth = np.array([[0, 0, 6, 3]])
assert (bboxes == truth).all()
def test_bitmap_mask_flip():
# flip with empty bitmap masks
raw_masks = dummy_raw_bitmap_masks((0, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
flipped_masks = bitmap_masks.flip(flip_direction='horizontal')
assert len(flipped_masks) == 0
assert flipped_masks.height == 28
assert flipped_masks.width == 28
# horizontally flip with bitmap masks contain 3 instances
raw_masks = dummy_raw_bitmap_masks((3, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
flipped_masks = bitmap_masks.flip(flip_direction='horizontal')
flipped_flipped_masks = flipped_masks.flip(flip_direction='horizontal')
assert flipped_masks.masks.shape == (3, 28, 28)
assert (bitmap_masks.masks == flipped_flipped_masks.masks).all()
assert (flipped_masks.masks == raw_masks[:, :, ::-1]).all()
# vertically flip with bitmap masks contain 3 instances
raw_masks = dummy_raw_bitmap_masks((3, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
flipped_masks = bitmap_masks.flip(flip_direction='vertical')
flipped_flipped_masks = flipped_masks.flip(flip_direction='vertical')
assert len(flipped_masks) == 3
assert flipped_masks.height == 28
assert flipped_masks.width == 28
assert (bitmap_masks.masks == flipped_flipped_masks.masks).all()
assert (flipped_masks.masks == raw_masks[:, ::-1, :]).all()
def test_bitmap_mask_resize():
# resize with empty bitmap masks
raw_masks = dummy_raw_bitmap_masks((0, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
resized_masks = bitmap_masks.resize((56, 72))
assert len(resized_masks) == 0
assert resized_masks.height == 56
assert resized_masks.width == 72
# resize with bitmap masks contain 1 instances
raw_masks = np.diag(np.ones(4, dtype=np.uint8))[np.newaxis, ...]
bitmap_masks = BitmapMasks(raw_masks, 4, 4)
resized_masks = bitmap_masks.resize((8, 8))
assert len(resized_masks) == 1
assert resized_masks.height == 8
assert resized_masks.width == 8
truth = np.array([[[1, 1, 0, 0, 0, 0, 0, 0], [1, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 1, 0, 0, 0, 0], [0, 0, 1, 1, 0, 0, 0, 0],
[0, 0, 0, 0, 1, 1, 0, 0], [0, 0, 0, 0, 1, 1, 0, 0],
[0, 0, 0, 0, 0, 0, 1, 1], [0, 0, 0, 0, 0, 0, 1, 1]]])
assert (resized_masks.masks == truth).all()
# resize to non-square
raw_masks = np.diag(np.ones(4, dtype=np.uint8))[np.newaxis, ...]
bitmap_masks = BitmapMasks(raw_masks, 4, 4)
resized_masks = bitmap_masks.resize((4, 8))
assert len(resized_masks) == 1
assert resized_masks.height == 4
assert resized_masks.width == 8
truth = np.array([[[1, 1, 0, 0, 0, 0, 0, 0], [0, 0, 1, 1, 0, 0, 0, 0],
[0, 0, 0, 0, 1, 1, 0, 0], [0, 0, 0, 0, 0, 0, 1, 1]]])
assert (resized_masks.masks == truth).all()
def test_bitmap_mask_area():
# area of empty bitmap mask
raw_masks = dummy_raw_bitmap_masks((0, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
assert bitmap_masks.areas.sum() == 0
# area of bitmap masks contain 3 instances
raw_masks = dummy_raw_bitmap_masks((3, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
areas = bitmap_masks.areas
assert len(areas) == 3
assert (areas == raw_masks.sum((1, 2))).all()
def dummy_masks(h, w, num_obj=3, mode='bitmap'):
assert mode in ('polygon', 'bitmap')
if mode == 'bitmap':
masks = np.random.randint(0, 2, (num_obj, h, w), dtype=np.uint8)
masks = BitmapMasks(masks, h, w)
else:
masks = []
for i in range(num_obj):
masks.append([])
masks[-1].append(
np.random.uniform(0, min(h - 1, w - 1), (8 + 4 * i, )))
masks[-1].append(
np.random.uniform(0, min(h - 1, w - 1), (10 + 4 * i, )))
masks = PolygonMasks(masks, h, w)
return masks
def test_mask_head_loss():
"""Test mask head loss when mask target is empty."""
self = FCNMaskHead(num_convs=1,
roi_feat_size=6,
in_channels=8,
conv_out_channels=8,
num_classes=8)
# Dummy proposals
proposal_list = [
torch.Tensor([[23.6667, 23.8757, 228.6326, 153.8874]]),
]
gt_bboxes = [
torch.Tensor([[23.6667, 23.8757, 238.6326, 151.8874]]),
]
gt_labels = [torch.LongTensor([2])]
sampling_results = _dummy_bbox_sampling(proposal_list, gt_bboxes,
gt_labels)
# create dummy mask
import numpy as np
from mmdet.core import BitmapMasks
dummy_mask = np.random.randint(0, 2, (1, 160, 240), dtype=np.uint8)
gt_masks = [BitmapMasks(dummy_mask, 160, 240)]
# create dummy train_cfg
train_cfg = mmcv.Config(dict(mask_size=12, mask_thr_binary=0.5))
# Create dummy features "extracted" for each sampled bbox
num_sampled = sum(len(res.bboxes) for res in sampling_results)
dummy_feats = torch.rand(num_sampled, 8, 6, 6)
mask_pred = self.forward(dummy_feats)
mask_targets = self.get_targets(sampling_results, gt_masks, train_cfg)
pos_labels = torch.cat([res.pos_gt_labels for res in sampling_results])
loss_mask = self.loss(mask_pred, mask_targets, pos_labels)
onegt_mask_loss = sum(loss_mask['loss_mask'])
assert onegt_mask_loss.item() > 0, 'mask loss should be non-zero'
# test mask_iou_head
mask_iou_head = MaskIoUHead(num_convs=1,
num_fcs=1,
roi_feat_size=6,
in_channels=8,
conv_out_channels=8,
fc_out_channels=8,
num_classes=8)
pos_mask_pred = mask_pred[range(mask_pred.size(0)), pos_labels]
mask_iou_pred = mask_iou_head(dummy_feats, pos_mask_pred)
pos_mask_iou_pred = mask_iou_pred[range(mask_iou_pred.size(0)), pos_labels]
mask_iou_targets = mask_iou_head.get_targets(sampling_results, gt_masks,
pos_mask_pred, mask_targets,
train_cfg)
loss_mask_iou = mask_iou_head.loss(pos_mask_iou_pred, mask_iou_targets)
onegt_mask_iou_loss = loss_mask_iou['loss_mask_iou'].sum()
assert onegt_mask_iou_loss.item() >= 0
def _load_side_face_map(self, results):
"""loading side face map
Args:
results (dict): Result dict from :obj:`dataset`.
Returns:
dict: The dict contains loaded side face map annotations.
"""
if self.file_client is None:
self.file_client = mmcv.FileClient(**self.file_client_args)
filename = osp.join(results['side_face_prefix'],
results['ann_info']['side_face_map'])
img_bytes = self.file_client.get(filename)
side_face_maps = mmcv.imfrombytes(img_bytes,
flag='unchanged').squeeze()
h, w = results['img_info']['height'], results['img_info']['width']
mask_num = len(results['ann_info']['masks'])
gt_side_face_maps = BitmapMasks(
[side_face_maps for _ in range(mask_num)], h, w)
results['gt_side_face_maps'] = gt_side_face_maps
results['side_face_fields'].append('gt_side_face_maps')
return results
def _masks(record):
if len(record.detect.masks) == 0:
raise RuntimeError("Negative samples still needs to be implemented")
else:
mask = record.detect.masks.data
_, h, w = mask.shape
return BitmapMasks(mask, height=h, width=w)
def _masks(record):
if len(record.detection.masks) == 0:
raise RuntimeError(
"Negative samples still needs to be implemented. This error might appear due to cropping, please check your transformations."
)
else:
mask = record.detection.mask_array.data
_, h, w = mask.shape
return BitmapMasks(mask, height=h, width=w)
def generate_targets(self, results):
"""Generate the gt targets for DBNet.
Args:
results (dict): The input result dictionary.
Returns:
results (dict): The output result dictionary.
"""
assert isinstance(results, dict)
if 'bbox_fields' in results:
results['bbox_fields'].clear()
ignore_tags = self.find_invalid(results)
results, ignore_tags = self.ignore_texts(results, ignore_tags)
h, w, _ = results['img_shape']
polygons = results['gt_masks'].masks
# generate gt_shrink_kernel
gt_shrink, ignore_tags = self.generate_kernels((h, w),
polygons,
self.shrink_ratio,
ignore_tags=ignore_tags)
results, ignore_tags = self.ignore_texts(results, ignore_tags)
# genenrate gt_shrink_mask
polygons_ignore = results['gt_masks_ignore'].masks
gt_shrink_mask = self.generate_effective_mask((h, w), polygons_ignore)
# generate gt_threshold and gt_threshold_mask
polygons = results['gt_masks'].masks
gt_thr, gt_thr_mask = self.generate_thr_map((h, w), polygons)
results['mask_fields'].clear() # rm gt_masks encoded by polygons
results.pop('gt_labels', None)
results.pop('gt_masks', None)
results.pop('gt_bboxes', None)
results.pop('gt_bboxes_ignore', None)
mapping = {
'gt_shrink': gt_shrink,
'gt_shrink_mask': gt_shrink_mask,
'gt_thr': gt_thr,
'gt_thr_mask': gt_thr_mask
}
for key, value in mapping.items():
value = value if isinstance(value, list) else [value]
results[key] = BitmapMasks(value, h, w)
results['mask_fields'].append(key)
return results
def bbox2mask(bboxes, w, h, mask_type='polygon'):
polys = bt.bbox2type(bboxes, 'poly')
assert mask_type in ['polygon', 'bitmap']
if mask_type == 'bitmap':
masks = []
for poly in polys:
rles = maskUtils.frPyObjects([poly.tolist()], h, w)
masks.append(maskUtils.decode(rles[0]))
gt_masks = BitmapMasks(masks, h, w)
else:
gt_masks = PolygonMasks([[poly] for poly in polys], h, w)
return gt_masks
def _load_masks(self, results):
h, w = results['img_info']['height'], results['img_info']['width']
gt_masks = results['ann_info']['masks']
if self.poly2mask:
gt_masks = BitmapMasks(
[self._poly2mask(mask, h, w) for mask in gt_masks], h, w)
else:
gt_masks = PolygonMasks(
[self.process_polygons(polygons) for polygons in gt_masks], h,
w)
results['gt_masks'] = gt_masks
results['mask_fields'].append('gt_masks')
return results
def test_panloss():
panloss = losses.PANLoss()
# test bitmasks2tensor
mask = [[1, 0, 1], [1, 1, 1], [0, 0, 1]]
target = [[1, 0, 1, 0, 0], [1, 1, 1, 0, 0], [0, 0, 1, 0, 0],
[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]
masks = [np.array(mask)]
bitmasks = BitmapMasks(masks, 3, 3)
target_sz = (6, 5)
results = panloss.bitmasks2tensor([bitmasks], target_sz)
assert len(results) == 1
assert torch.sum(torch.abs(results[0].float() -
torch.Tensor(target))).item() == 0
def __call__(self, results):
gt_mask = results[self.instance_key]
mask = None
if len(gt_mask.masks) > 0:
mask = self.generate_mask(gt_mask, self.mask_type)
results['crop_offset'] = self.sample_offset(mask,
results['img'].shape[:2])
# crop img. bbox = [x1,y1,x2,y2]
img, bbox = self.crop_img(results['img'], results['crop_offset'],
self.target_size)
results['img'] = img
img_shape = img.shape
results['img_shape'] = img_shape
# crop masks
for key in results.get('mask_fields', []):
results[key] = results[key].crop(bbox)
# for mask rcnn
for key in results.get('bbox_fields', []):
results[key], kept_inx = self.crop_bboxes(results[key], bbox)
if key == 'gt_bboxes':
# ignore gt_labels accordingly
if 'gt_labels' in results:
ori_labels = results['gt_labels']
ori_inst_num = len(ori_labels)
results['gt_labels'] = [
ori_labels[idx] for idx in range(ori_inst_num)
if idx in kept_inx
]
# ignore g_masks accordingly
if 'gt_masks' in results:
ori_mask = results['gt_masks'].masks
kept_mask = [
ori_mask[idx] for idx in range(ori_inst_num)
if idx in kept_inx
]
target_h, target_w = bbox[3] - bbox[1], bbox[2] - bbox[0]
if len(kept_inx) > 0:
kept_mask = np.stack(kept_mask)
else:
kept_mask = np.empty((0, target_h, target_w),
dtype=np.float32)
results['gt_masks'] = BitmapMasks(kept_mask, target_h,
target_w)
return results
def _load_masks(self, results):
ann_info = results['ann_info']
h, w = results['img_info']['height'], results['img_info']['width']
if self.use_img_shape:
if results.get('ori_shape', None):
h, w = results['ori_shape'][:2]
results['img_info']['height'] = h
results['img_info']['width'] = w
else:
warnings.warn('"ori_shape" not in results, use the shape '
'in "img_info" instead.')
gt_masks = ann_info['masks']
if self.poly2mask:
gt_masks = BitmapMasks(
[self._poly2mask(mask, h, w) for mask in gt_masks], h, w)
else:
gt_masks = PolygonMasks(
[self.process_polygons(polygons) for polygons in gt_masks], h,
w)
gt_masks_ignore = ann_info.get('masks_ignore', None)
if gt_masks_ignore is not None:
if self.poly2mask:
gt_masks_ignore = BitmapMasks(
[self._poly2mask(mask, h, w) for mask in gt_masks_ignore],
h, w)
else:
gt_masks_ignore = PolygonMasks([
self.process_polygons(polygons)
for polygons in gt_masks_ignore
], h, w)
results['gt_masks_ignore'] = gt_masks_ignore
results['mask_fields'].append('gt_masks_ignore')
results['gt_masks'] = gt_masks
results['mask_fields'].append('gt_masks')
return results
def test_random_rotate(mock_random):
mock_random.side_effect = [0.5, 0]
results = {}
img = np.random.rand(5, 5)
results['img'] = img.copy()
results['mask_fields'] = ['masks']
gt_kernels = [results['img'].copy()]
results['masks'] = BitmapMasks(gt_kernels, 5, 5)
rotater = transforms.RandomRotateTextDet()
results = rotater(results)
assert np.allclose(results['img'], img)
assert np.allclose(results['masks'].masks, img)
def _load_masks_and_semantic_segs(self, results):
"""Private function to load mask and semantic segmentation annotations.
In gt_semantic_seg, the foreground label is from `0` to
`num_things - 1`, the background label is from `num_things` to
`num_things + num_stuff - 1`, 255 means the ignored label (`VOID`).
Args:
results (dict): Result dict from :obj:`mmdet.CustomDataset`.
Returns:
dict: The dict contains loaded mask and semantic segmentation
annotations. `BitmapMasks` is used for mask annotations.
"""
if self.file_client is None:
self.file_client = mmcv.FileClient(**self.file_client_args)
filename = osp.join(results['seg_prefix'],
results['ann_info']['seg_map'])
img_bytes = self.file_client.get(filename)
pan_png = mmcv.imfrombytes(img_bytes,
flag='color',
channel_order='rgb').squeeze()
pan_png = rgb2id(pan_png)
gt_masks = []
gt_seg = np.zeros_like(pan_png) + 255 # 255 as ignore
for mask_info in results['ann_info']['masks']:
mask = (pan_png == mask_info['id'])
gt_seg = np.where(mask, mask_info['category'], gt_seg)
# The legal thing masks
if mask_info.get('is_thing'):
gt_masks.append(mask.astype(np.uint8))
if self.with_mask:
h, w = results['img_info']['height'], results['img_info']['width']
gt_masks = BitmapMasks(gt_masks, h, w)
results['gt_masks'] = gt_masks
results['mask_fields'].append('gt_masks')
if self.with_seg:
results['gt_semantic_seg'] = gt_seg
results['seg_fields'].append('gt_semantic_seg')
return results
def __call__(self, results):
img_shape = results['img_shape']
resize_shape = results['resize_shape']
h_scale = 1.0 * resize_shape[0] / img_shape[0]
w_scale = 1.0 * resize_shape[1] / img_shape[1]
char_boxes, char_inds = [], []
char_num = len(results['ann_info'][self.box_type])
for i in range(char_num):
char_box = results['ann_info'][self.box_type][i]
num_points = 2 if self.box_type == 'char_rects' else 4
for j in range(num_points):
char_box[j * 2] = round(char_box[j * 2] * w_scale)
char_box[j * 2 + 1] = round(char_box[j * 2 + 1] * h_scale)
char_boxes.append(char_box)
char = results['ann_info']['chars'][i]
char_ind = self.label_convertor.str2idx([char])[0][0]
char_inds.append(char_ind)
resize_shape = tuple(results['resize_shape'][:2])
pad_shape = tuple(results['pad_shape'][:2])
binary_target = self.generate_kernels(
resize_shape,
pad_shape,
char_boxes,
char_inds,
shrink_ratio=self.attn_shrink_ratio,
binary=True)
seg_target = self.generate_kernels(
resize_shape,
pad_shape,
char_boxes,
char_inds,
shrink_ratio=self.seg_shrink_ratio,
binary=False)
mask = np.ones(pad_shape, dtype=np.int32)
mask[:resize_shape[0], resize_shape[1]:] = 0
results['gt_kernels'] = BitmapMasks([binary_target, seg_target, mask],
pad_shape[0], pad_shape[1])
results['mask_fields'] = ['gt_kernels']
return results
def generate_targets(self, results):
"""Generate the gt targets for DRRG.
Args:
results (dict): The input result dictionary.
Returns:
results (dict): The output result dictionary.
"""
assert isinstance(results, dict)
polygon_masks = results['gt_masks'].masks
polygon_masks_ignore = results['gt_masks_ignore'].masks
h, w, _ = results['img_shape']
gt_text_mask = self.generate_text_region_mask((h, w), polygon_masks)
gt_mask = self.generate_effective_mask((h, w), polygon_masks_ignore)
(center_lines, gt_center_region_mask, gt_top_height_map,
gt_bot_height_map, gt_sin_map,
gt_cos_map) = self.generate_center_mask_attrib_maps((h, w),
polygon_masks)
gt_comp_attribs = self.generate_comp_attribs(
center_lines, gt_text_mask, gt_center_region_mask,
gt_top_height_map, gt_bot_height_map, gt_sin_map, gt_cos_map)
results['mask_fields'].clear() # rm gt_masks encoded by polygons
mapping = {
'gt_text_mask': gt_text_mask,
'gt_center_region_mask': gt_center_region_mask,
'gt_mask': gt_mask,
'gt_top_height_map': gt_top_height_map,
'gt_bot_height_map': gt_bot_height_map,
'gt_sin_map': gt_sin_map,
'gt_cos_map': gt_cos_map
}
for key, value in mapping.items():
value = value if isinstance(value, list) else [value]
results[key] = BitmapMasks(value, h, w)
results['mask_fields'].append(key)
results['gt_comp_attribs'] = gt_comp_attribs
return results
def __call__(self, results):
if np.random.random_sample() < self.rotate_ratio:
# rotate imgs
results['rotated_angle'] = self.sample_angle(self.max_angle)
img = self.rotate_img(results['img'], results['rotated_angle'])
results['img'] = img
img_shape = img.shape
results['img_shape'] = img_shape
# rotate masks
for key in results.get('mask_fields', []):
masks = results[key].masks
mask_list = []
for m in masks:
rotated_m = self.rotate_img(m, results['rotated_angle'])
mask_list.append(rotated_m)
results[key] = BitmapMasks(mask_list, *(img_shape[:2]))
return results
def generate_targets(self, results):
"""Generate the gt targets for PANet.
Args:
results (dict): The input result dictionary.
Returns:
results (dict): The output result dictionary.
"""
assert isinstance(results, dict)
polygon_masks = results['gt_masks'].masks
polygon_masks_ignore = results['gt_masks_ignore'].masks
h, w, _ = results['img_shape']
gt_kernels = []
for ratio in self.shrink_ratio:
mask, _ = self.generate_kernels((h, w),
polygon_masks,
ratio,
max_shrink=self.max_shrink,
ignore_tags=None)
gt_kernels.append(mask)
gt_mask = self.generate_effective_mask((h, w), polygon_masks_ignore)
results['mask_fields'].clear() # rm gt_masks encoded by polygons
if 'bbox_fields' in results:
results['bbox_fields'].clear()
results.pop('gt_labels', None)
results.pop('gt_masks', None)
results.pop('gt_bboxes', None)
results.pop('gt_bboxes_ignore', None)
mapping = {'gt_kernels': gt_kernels, 'gt_mask': gt_mask}
for key, value in mapping.items():
value = value if isinstance(value, list) else [value]
results[key] = BitmapMasks(value, h, w)
results['mask_fields'].append(key)
return results
def test_bitmap_mask_to_ndarray():
# empty bitmap masks to ndarray
raw_masks = dummy_raw_bitmap_masks((0, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
ndarray_masks = bitmap_masks.to_ndarray()
assert isinstance(ndarray_masks, np.ndarray)
assert ndarray_masks.shape == (0, 28, 28)
# bitmap masks contain 3 instances to ndarray
raw_masks = dummy_raw_bitmap_masks((3, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
ndarray_masks = bitmap_masks.to_ndarray()
assert isinstance(ndarray_masks, np.ndarray)
assert ndarray_masks.shape == (3, 28, 28)
assert (ndarray_masks == raw_masks).all()
def test_bitmap_mask_to_tensor():
# empty bitmap masks to tensor
raw_masks = dummy_raw_bitmap_masks((0, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
tensor_masks = bitmap_masks.to_tensor(dtype=torch.uint8, device='cpu')
assert isinstance(tensor_masks, torch.Tensor)
assert tensor_masks.shape == (0, 28, 28)
# bitmap masks contain 3 instances to tensor
raw_masks = dummy_raw_bitmap_masks((3, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
tensor_masks = bitmap_masks.to_tensor(dtype=torch.uint8, device='cpu')
assert isinstance(tensor_masks, torch.Tensor)
assert tensor_masks.shape == (3, 28, 28)
assert (tensor_masks.numpy() == raw_masks).all()
def switch_mask_type(masks, mtype='bitmap'):
if isinstance(masks, PolygonMasks) and mtype == 'bitmap':
width, height = masks.width, masks.height
bitmap_masks = []
for poly_per_obj in masks.masks:
rles = maskUtils.frPyObjects(poly_per_obj, height, width)
rle = maskUtils.merge(rles)
bitmap_masks.append(maskUtils.decode(rle).astype(np.uint8))
masks = BitmapMasks(bitmap_masks, height, width)
elif isinstance(masks, BitmapMasks) and mtype == 'polygon':
width, height = masks.width, masks.height
polygons = []
for bitmask in masks.masks:
try:
contours, _ = cv2.findContours(bitmask, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
except ValueError:
_, contours, _ = cv2.findContours(bitmask, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
polygons.append(list(contours))
masks = PolygonMasks(polygons, width, height)
return masks
def test_bitmap_mask_crop():
# crop with empty bitmap masks
dummy_bbox = np.array([0, 10, 10, 27], dtype=np.int)
raw_masks = dummy_raw_bitmap_masks((0, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
cropped_masks = bitmap_masks.crop(dummy_bbox)
assert len(cropped_masks) == 0
assert cropped_masks.height == 17
assert cropped_masks.width == 10
# crop with bitmap masks contain 3 instances
raw_masks = dummy_raw_bitmap_masks((3, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
cropped_masks = bitmap_masks.crop(dummy_bbox)
assert len(cropped_masks) == 3
assert cropped_masks.height == 17
assert cropped_masks.width == 10
x1, y1, x2, y2 = dummy_bbox
assert (cropped_masks.masks == raw_masks[:, y1:y2, x1:x2]).all()
# crop with invalid bbox
with pytest.raises(AssertionError):
dummy_bbox = dummy_bboxes(2, 28, 28)
bitmap_masks.crop(dummy_bbox)
def test_bitmap_mask_pad():
# pad with empty bitmap masks
raw_masks = dummy_raw_bitmap_masks((0, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
padded_masks = bitmap_masks.pad((56, 56))
assert len(padded_masks) == 0
assert padded_masks.height == 56
assert padded_masks.width == 56
# pad with bitmap masks contain 3 instances
raw_masks = dummy_raw_bitmap_masks((3, 28, 28))
bitmap_masks = BitmapMasks(raw_masks, 28, 28)
padded_masks = bitmap_masks.pad((56, 56))
assert len(padded_masks) == 3
assert padded_masks.height == 56
assert padded_masks.width == 56
assert (padded_masks.masks[:, 28:, 28:] == 0).all()
