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 convert_bbox(bbox, img_h, img_w):
rle = mask_utils.frPyObjects([bbox.as_polygon()], img_h, img_w)[0]
return RleMask(rle=rle,
label=bbox.label,
z_order=bbox.z_order,
id=bbox.id,
attributes=bbox.attributes,
group=bbox.group)
def convert_polygon(polygon, img_h, img_w):
rle = mask_utils.frPyObjects([polygon.points], img_h, img_w)[0]
return RleMask(rle=rle,
label=polygon.label,
z_order=polygon.z_order,
id=polygon.id,
attributes=polygon.attributes,
group=polygon.group)
def _load_annotations(item):
parsed = item['annotations']
loaded = []
for ann in parsed:
ann_id = ann.get('id')
ann_type = AnnotationType[ann['type']]
attributes = ann.get('attributes')
group = ann.get('group')
label_id = ann.get('label_id')
z_order = ann.get('z_order')
points = ann.get('points')
if ann_type == AnnotationType.label:
loaded.append(Label(label=label_id,
id=ann_id, attributes=attributes, group=group))
elif ann_type == AnnotationType.mask:
rle = ann['rle']
rle['counts'] = rle['counts'].encode('ascii')
loaded.append(RleMask(rle=rle, label=label_id,
id=ann_id, attributes=attributes, group=group,
z_order=z_order))
elif ann_type == AnnotationType.polyline:
loaded.append(PolyLine(points, label=label_id,
id=ann_id, attributes=attributes, group=group,
z_order=z_order))
elif ann_type == AnnotationType.polygon:
loaded.append(Polygon(points, label=label_id,
id=ann_id, attributes=attributes, group=group,
z_order=z_order))
elif ann_type == AnnotationType.bbox:
x, y, w, h = ann['bbox']
loaded.append(Bbox(x, y, w, h, label=label_id,
id=ann_id, attributes=attributes, group=group,
z_order=z_order))
elif ann_type == AnnotationType.points:
loaded.append(Points(points, label=label_id,
id=ann_id, attributes=attributes, group=group,
z_order=z_order))
elif ann_type == AnnotationType.caption:
caption = ann.get('caption')
loaded.append(Caption(caption,
id=ann_id, attributes=attributes, group=group))
else:
raise NotImplementedError()
return loaded
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
def _load_annotations(self, ann, image_info=None):
parsed_annotations = []
ann_id = ann.get('id')
attributes = {}
if 'score' in ann:
attributes['score'] = ann['score']
group = ann_id # make sure all tasks' annotations are merged
if self._task in [CocoTask.instances, CocoTask.person_keypoints]:
x, y, w, h = ann['bbox']
label_id = self._get_label_id(ann)
is_crowd = bool(ann['iscrowd'])
attributes['is_crowd'] = is_crowd
if self._task is CocoTask.person_keypoints:
keypoints = ann['keypoints']
points = [p for i, p in enumerate(keypoints) if i % 3 != 2]
visibility = keypoints[2::3]
parsed_annotations.append(
Points(points,
visibility,
label=label_id,
id=ann_id,
attributes=attributes,
group=group))
segmentation = ann.get('segmentation')
if segmentation and segmentation != [[]]:
rle = None
if isinstance(segmentation, list):
if not self._merge_instance_polygons:
# polygon - a single object can consist of multiple parts
for polygon_points in segmentation:
parsed_annotations.append(
Polygon(points=polygon_points,
label=label_id,
id=ann_id,
attributes=attributes,
group=group))
else:
# merge all parts into a single mask RLE
img_h = image_info['height']
img_w = image_info['width']
rles = mask_utils.frPyObjects(segmentation, img_h,
img_w)
rle = mask_utils.merge(rles)
elif isinstance(segmentation['counts'], list):
# uncompressed RLE
img_h = image_info['height']
img_w = image_info['width']
mask_h, mask_w = segmentation['size']
if img_h == mask_h and img_w == mask_w:
rle = mask_utils.frPyObjects([segmentation], mask_h,
mask_w)[0]
else:
log.warning(
"item #%s: mask #%s "
"does not match image size: %s vs. %s. "
"Skipping this annotation.", image_info['id'],
ann_id, (mask_h, mask_w), (img_h, img_w))
else:
# compressed RLE
rle = segmentation
if rle is not None:
parsed_annotations.append(
RleMask(rle=rle,
label=label_id,
id=ann_id,
attributes=attributes,
group=group))
else:
parsed_annotations.append(
Bbox(x,
y,
w,
h,
label=label_id,
id=ann_id,
attributes=attributes,
group=group))
elif self._task is CocoTask.labels:
label_id = self._get_label_id(ann)
parsed_annotations.append(
Label(label=label_id,
id=ann_id,
attributes=attributes,
group=group))
elif self._task is CocoTask.captions:
caption = ann['caption']
parsed_annotations.append(
Caption(caption, id=ann_id, attributes=attributes,
group=group))
else:
raise NotImplementedError()
return parsed_annotations