def combine_stuff_masks(
rles: List[RLE],
class_ids: List[int],
inst_ids: List[int],
classes: List[Category],
) -> Tuple[List[RLE], List[int], List[int]]:
"""For each stuff class, combine masks into a single mask."""
combine_rles: List[RLE] = []
combine_cids: List[int] = []
combine_iids: List[int] = []
for class_id in sorted(set(class_ids)):
category = classes[class_id]
rles_c = [
rle for rle, c_id in zip(rles, class_ids) if c_id == class_id
]
iids_c = [
iid for iid, c_id in zip(inst_ids, class_ids) if c_id == class_id
]
if category.isThing is None or category.isThing:
combine_rles.extend(rles_c)
combine_cids.extend([class_id] * len(rles_c))
combine_iids.extend(iids_c)
else:
combine_mask: NDArrayU8 = sum( # type: ignore
rle_to_mask(rle) for rle in rles_c)
combine_rles.append(mask_to_rle(combine_mask))
combine_cids.append(class_id)
combine_iids.append(iids_c[0])
return combine_rles, combine_cids, combine_iids
def insseg_to_rle(frame: Frame, input_dir: str,
categories: List[Category]) -> Frame:
"""Convert ins_seg bitmasks to rle."""
ann_score: Dict[str, List[Tuple[int, float]]] = {}
img_name = frame.name.replace(".jpg", ".png")
ann_score[img_name] = []
bitmask: NDArrayU8 = np.array(
Image.open(os.path.join(input_dir, img_name)),
dtype=np.uint8,
)
if frame.labels is None:
return frame
for label in frame.labels:
assert label.index is not None
assert label.score is not None
ann_score[img_name].append((label.index, label.score))
masks, ann_ids, scores, category_ids = parse_res_bitmask(
ann_score[img_name], bitmask)
labels = []
for ann_id in ann_ids:
label = Label(
id=ann_id,
category=categories[category_ids[ann_id - 1] - 1].name,
score=scores[ann_id - 1],
)
label.rle = mask_to_rle((masks == ann_id).astype(np.uint8))
labels.append(label)
frame.labels = labels
return frame
def semseg_to_rle(frame: Frame, input_dir: str,
categories: List[Category]) -> Frame:
"""Convert sem_seg bitmasks to rle."""
frame.labels = []
img_name = frame.name.replace(".jpg", ".png")
bitmask: NDArrayU8 = np.array(
Image.open(os.path.join(input_dir, img_name)),
dtype=np.uint8,
)
category_ids: NDArrayU8 = np.unique(bitmask)
label_id = 0
for category_id in category_ids:
label = Label(id=str(label_id))
label.category = categories[category_id].name
label.rle = mask_to_rle((bitmask == category_id).astype(np.uint8))
frame.labels.append(label)
label_id += 1
return frame
def parse_seg_objects(
objects: List[Label],
classes: List[Category],
ignore_unknown_cats: bool = False,
image_size: Optional[ImageSize] = None,
) -> Tuple[List[RLEDict], NDArrayI32, NDArrayI32, List[RLEDict]]:
"""Parse segmentation objects under Scalabel formats."""
rles, labels, ids, ignore_rles = [], [], [], []
class_names = [category.name for category in classes]
for obj in objects:
if obj.rle is not None:
rle = obj.rle
elif obj.poly2d is not None:
assert (
image_size
is not None), "Requires ImageSize for Poly2D conversion to RLE"
rle = mask_to_rle(poly2ds_to_mask(image_size, obj.poly2d))
else:
continue
category = obj.category
if category in class_names:
if check_crowd(obj) or check_ignored(obj):
ignore_rles.append(rle)
else:
rles.append(rle)
labels.append(class_names.index(category))
ids.append(int(obj.id))
else:
if not ignore_unknown_cats:
raise KeyError(f"Unknown category: {category}")
if any(category.isThing is not None and not category.isThing
for category in classes):
rles, labels, ids = combine_stuff_masks(rles, labels, ids, classes)
rles_dict = [rle.dict() for rle in rles]
ignore_rles_dict = [rle.dict() for rle in ignore_rles]
labels_arr: NDArrayI32 = np.array(labels, dtype=np.int32)
ids_arr: NDArrayI32 = np.array(ids, dtype=np.int32)
return (rles_dict, labels_arr, ids_arr, ignore_rles_dict)
def segtrack_to_rle(frame: Frame, input_dir: str,
categories: List[Category]) -> Frame:
"""Convert seg_track bitmasks to rle."""
frame.labels = []
img_name = frame.name.replace(".jpg", ".png")
bitmask: NDArrayU8 = np.array(
Image.open(os.path.join(input_dir, img_name)),
dtype=np.uint8,
)
masks, instance_ids, _, category_ids = parse_bitmask(bitmask)
# video parameters
frame.name = frame.name.split("/")[-1]
frame.videoName = img_name.split("/")[0]
frame.frameIndex = int(img_name.split("-")[-1].split(".")[0]) - 1
for i, _ in enumerate(instance_ids):
label = Label(id=str(instance_ids[i]))
label.category = categories[category_ids[i] - 1].name
label.rle = mask_to_rle((masks == i + 1).astype(np.uint8))
frame.labels.append(label)
return frame