def crop_segments(cls, segment_anns, img_width, img_height):
segment_anns = sorted(segment_anns, key=lambda x: x.z_order)
segments = []
for s in segment_anns:
if s.type == AnnotationType.polygon:
segments.append(s.points)
elif s.type == AnnotationType.mask:
if isinstance(s, RleMask):
rle = s.rle
else:
rle = mask_tools.mask_to_rle(s.image)
segments.append(rle)
segments = mask_tools.crop_covered_segments(
segments, img_width, img_height)
new_anns = []
for ann, new_segment in zip(segment_anns, segments):
fields = {'z_order': ann.z_order, 'label': ann.label,
'id': ann.id, 'group': ann.group, 'attributes': ann.attributes
}
if ann.type == AnnotationType.polygon:
if fields['group'] is None:
fields['group'] = cls._make_group_id(
segment_anns + new_anns, fields['id'])
for polygon in new_segment:
new_anns.append(Polygon(points=polygon, **fields))
else:
rle = mask_tools.mask_to_rle(new_segment)
rle = mask_utils.frPyObjects(rle, *rle['size'])
new_anns.append(RleMask(rle=rle, **fields))
return new_anns
def crop_segments(cls, instances, img_width, img_height):
instances = sorted(instances, key=lambda x: x[0].z_order)
segment_map = []
segments = []
for inst_idx, (_, polygons, mask, _) in enumerate(instances):
if polygons:
segment_map.extend(inst_idx for p in polygons)
segments.extend(polygons)
elif mask is not None:
segment_map.append(inst_idx)
segments.append(mask)
segments = mask_tools.crop_covered_segments(segments, img_width,
img_height)
for inst_idx, inst in enumerate(instances):
new_segments = [
s for si_id, s in zip(segment_map, segments)
if si_id == inst_idx
]
if not new_segments:
inst[1] = []
inst[2] = None
continue
if inst[1]:
inst[1] = sum(new_segments, [])
else:
mask = mask_tools.merge_masks(new_segments)
inst[2] = mask_tools.mask_to_rle(mask)
return instances
def _to_rle(ann):
if ann.type == AnnotationType.polygon:
return mask_utils.frPyObjects([ann.points], h, w)
elif isinstance(ann, RleMask):
return [ann._rle]
elif ann.type == AnnotationType.mask:
return mask_utils.frPyObjects([mask_to_rle(ann.image)], h, w)
else:
raise TypeError("Unexpected arguments: %s, %s" % (a, b))
def _test_mask_to_rle(self, source_mask):
rle_uncompressed = mask_tools.mask_to_rle(source_mask)
from pycocotools import mask as mask_utils
resulting_mask = mask_utils.frPyObjects(
rle_uncompressed, *rle_uncompressed['size'])
resulting_mask = mask_utils.decode(resulting_mask)
self.assertTrue(np.array_equal(source_mask, resulting_mask),
'%s\n%s\n' % (source_mask, resulting_mask))
def test_can_crop_covered_segments(self):
image_size = [7, 7]
initial = [
[1, 1, 6, 1, 6, 6, 1, 6], # rectangle
mask_tools.mask_to_rle(
np.array([
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 1, 1, 0],
[0, 1, 1, 0, 1, 1, 0],
[0, 0, 0, 0, 0, 1, 0],
[0, 1, 1, 0, 0, 1, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 0, 0, 0, 0, 0, 0],
])),
[1, 1, 6, 6, 1, 6], # lower-left triangle
]
expected = [
np.array([
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 0],
[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, 0, 0],
]), # half-covered
np.array([
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 1, 1, 0],
[0, 0, 0, 0, 1, 1, 0],
[0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
]), # half-covered
mask_tools.rles_to_mask([initial[2]], *image_size), # unchanged
]
computed = mask_tools.crop_covered_segments(initial,
*image_size,
ratio_tolerance=0,
return_masks=True)
self.assertEqual(len(initial), len(computed))
for i, (e_mask, c_mask) in enumerate(zip(expected, computed)):
self.assertTrue(np.array_equal(e_mask, c_mask),
'#%s: %s\n%s\n' % (i, e_mask, c_mask))
def find_instance_parts(self, group, img_width, img_height):
boxes = [a for a in group if a.type == AnnotationType.bbox]
polygons = [a for a in group if a.type == AnnotationType.polygon]
masks = [a for a in group if a.type == AnnotationType.mask]
anns = boxes + polygons + masks
leader = anno_tools.find_group_leader(anns)
bbox = anno_tools.max_bbox(anns)
mask = None
polygons = [p.points for p in polygons]
if self._context._segmentation_mode == SegmentationMode.guess:
use_masks = True == leader.attributes.get(
'is_crowd',
find(masks, lambda x: x.label == leader.label) is not None)
elif self._context._segmentation_mode == SegmentationMode.polygons:
use_masks = False
elif self._context._segmentation_mode == SegmentationMode.mask:
use_masks = True
else:
raise NotImplementedError("Unexpected segmentation mode '%s'" % \
self._context._segmentation_mode)
if use_masks:
if polygons:
mask = mask_tools.rles_to_mask(polygons, img_width, img_height)
if masks:
masks = (m.image for m in masks)
if mask is not None:
masks += chain(masks, [mask])
mask = mask_tools.merge_masks(masks)
if mask is not None:
mask = mask_tools.mask_to_rle(mask)
polygons = []
else:
if masks:
mask = mask_tools.merge_masks(m.image for m in masks)
polygons += mask_tools.mask_to_polygons(mask)
mask = None
return [leader, polygons, mask, bbox]
def merge_segments(cls,
instance,
img_width,
img_height,
include_polygons=False):
polygons = [a for a in instance if a.type == AnnotationType.polygon]
masks = [a for a in instance if a.type == AnnotationType.mask]
if not polygons and not masks:
return []
if not polygons and len(masks) == 1:
return masks
leader = find_group_leader(polygons + masks)
instance = []
# Build the resulting mask
mask = None
if include_polygons and polygons:
polygons = [p.points for p in polygons]
mask = mask_tools.rles_to_mask(polygons, img_width, img_height)
else:
instance += polygons # keep unused polygons
if masks:
masks = (m.image for m in masks)
if mask is not None:
masks = chain(masks, [mask])
mask = mask_tools.merge_masks(masks)
if mask is None:
return instance
mask = mask_tools.mask_to_rle(mask)
mask = mask_utils.frPyObjects(mask, *mask['size'])
instance.append(
RleMask(rle=mask,
label=leader.label,
z_order=leader.z_order,
id=leader.id,
attributes=leader.attributes,
group=leader.group))
return instance
