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 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