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 deal_bdd100k_category(label: Label, bdd100k_config: BDD100KConfig,
cat_name2id: Dict[str, int]) -> Optional[Label]:
"""Deal with BDD100K category."""
category_name = label.category
if (bdd100k_config.name_mapping is not None
and category_name in bdd100k_config.name_mapping):
category_name = bdd100k_config.name_mapping[category_name]
if category_name not in cat_name2id:
if bdd100k_config.remove_ignored:
result = None
elif bdd100k_config.ignored_as_class:
assert IGNORED in cat_name2id
category_name = IGNORED
label.category = category_name
result = label
else:
assert bdd100k_config.ignored_mapping is not None
assert category_name in bdd100k_config.ignored_mapping
category_name = bdd100k_config.ignored_mapping[category_name]
if label.attributes is None:
label.attributes = {}
label.attributes[IGNORED] = True
label.category = category_name
result = label
else:
label.category = category_name
result = label
return result
def box_track_to_bdd100k(dataset, results, out_base, nproc):
bdd100k = []
track_base = osp.join(out_base, "box_track")
if not osp.exists(track_base):
os.makedirs(track_base)
print(f'\nStart converting to BDD100K box tracking format')
for idx, track_dict in tqdm(enumerate(results['track_results'])):
img_name = dataset.data_infos[idx]['file_name']
frame_index = dataset.data_infos[idx]['frame_id']
vid_name = os.path.split(img_name)[0]
frame = Frame(
name=img_name,
video_name=vid_name,
frame_index=frame_index,
labels=[])
for id_, instance in track_dict.items():
bbox = instance['bbox']
cls_ = instance['label']
label = Label(
id=id_,
score=bbox[-1],
box2d=bbox_to_box2d(xyxy2xywh(bbox)),
category=CATEGORIES[cls_ + 1])
frame.labels.append(label)
bdd100k.append(frame)
print(f'\nWriting the converted json')
out_path = osp.join(out_base, "box_track.json")
save(out_path, bdd100k)
def test_set_instance_color(self) -> None:
"""Check color setting."""
label = Label(
id="tmp",
attributes=dict(truncated=True, crowd=False),
)
color = set_instance_color(label, 15, 300, False)
gt_color = np.array([15, 8, 1, 44])
self.assertTrue((color == gt_color).all())
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
def det_to_bdd(self, results, out_base, nproc):
bdd100k = []
ann_id = 0
print(f'\nStart converting to BDD100K detection format')
for idx, bboxes_list in tqdm(enumerate(results['bbox_result'])):
img_name = self.data_infos[idx]['file_name']
frame = Frame(name=img_name, labels=[])
for cls_, bboxes in enumerate(bboxes_list):
for bbox in bboxes:
ann_id += 1
label = Label(id=ann_id,
score=bbox[-1],
box2d=bbox_to_box2d(xyxy2xywh(bbox)),
category=CATEGORIES[cls_ + 1])
frame.labels.append(label)
bdd100k.append(frame)
print(f'\nWriting the converted json')
out_path = osp.join(out_base, "det.json")
save(out_path, bdd100k)
def ins_seg_to_bdd100k(dataset, results, out_base, nproc=4):
bdd100k = []
bitmask_base = osp.join(out_base, "ins_seg")
if not osp.exists(bitmask_base):
os.makedirs(bitmask_base)
if 'bbox_results' in results and 'segm_results' in results:
results = [[bbox, segm] for bbox, segm in zip(results['bbox_results'],
results['segm_results'])]
track_dicts = []
img_names = [
dataset.data_infos[idx]['file_name'] for idx in range(len(results))
]
print(f'\nStart converting to BDD100K instance segmentation format')
ann_id = 0
for idx, [bboxes_list, segms_list] in enumerate(results):
index = 0
frame = Frame(name=img_names[idx], labels=[])
track_dict = {}
for cls_, (bboxes, segms) in enumerate(zip(bboxes_list, segms_list)):
for bbox, segm in zip(bboxes, segms):
ann_id += 1
index += 1
label = Label(id=str(ann_id), index=index, score=bbox[-1])
frame.labels.append(label)
instance = {'bbox': bbox, 'segm': segm, 'label': cls_}
track_dict[index] = instance
bdd100k.append(frame)
track_dicts.append(track_dict)
print(f'\nWriting the converted json')
out_path = osp.join(out_base, 'ins_seg.json')
save(out_path, bdd100k)
mask_merge_parallel(track_dicts, img_names, bitmask_base, nproc)