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 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 main() -> None:
"""Main."""
args = parse_args()
assert os.path.isdir(args.input)
dataset = load(args.input, args.nproc)
if args.config is not None:
bdd100k_config = load_bdd100k_config(args.config)
elif dataset.config is not None:
bdd100k_config = BDD100KConfig(config=dataset.config)
else:
bdd100k_config = load_bdd100k_config(args.mode)
categories = get_leaf_categories(bdd100k_config.scalabel.categories)
convert_funcs: Dict[str, ToRLEFunc] = dict(
ins_seg=insseg_to_rle,
sem_seg=semseg_to_rle,
drivable=semseg_to_rle,
seg_track=segtrack_to_rle,
)
if args.mode == "ins_seg":
assert args.score_file is not None
frames = load(args.score_file).frames
assert all(
os.path.exists(
os.path.join(args.input, frame.name.replace(".jpg", ".png")))
for frame in frames), "Missing some bitmasks."
elif args.mode in ("sem_seg", "drivable", "seg_track"):
files = list_files(args.input)
frames = []
for file in files:
if not file.endswith(".png") and not file.endswith(".jpg"):
continue
frame = Frame(name=file.replace(".png", ".jpg"), labels=[])
frames.append(frame)
else:
return
if args.nproc > 1:
with Pool(args.nproc) as pool:
frames = pool.map(
partial(
convert_funcs[args.mode],
input_dir=args.input,
categories=categories,
),
tqdm(frames),
)
else:
frames = [
convert_funcs[args.mode](frame, args.input, categories)
for frame in tqdm(frames)
]
save(args.output, frames)
def mock_load(*_) -> Dataset: # type: ignore
"""Mock load to return a dummy dataset with one frame."""
return Dataset(frames=[
Frame(
name="7d06fefd-f7be05a6.png",
labels=[],
videoName="7d06fefd-f7be05a6",
frameIndex=0,
)
])
def test_frame_load() -> None:
"""Test loading frame."""
json = (
'{"name": 1, "videoName": "a", "size": [10, 20], '
'"labels":[{"id": 1, "box2d": {"x1": 1, "y1": 2, "x2": 3, "y2": 4}}]}')
# pylint and mypy doesn't recognize from_json method
frame = Frame.from_json(json) # type: ignore # pylint: disable=no-member
assert frame.name == "1"
assert frame.video_name == "a"
assert frame.labels[0].id == "1"
def test_semseg_to_rle(self) -> None:
"""Test sem_seg to rle conversion."""
mask_dir = "./testcases/to_rle/sem_seg/masks"
mask = "0.jpg"
categories = get_leaf_categories(
load_bdd100k_config("sem_seg").scalabel.categories
)
frame = Frame(name=mask)
new_frame = semseg_to_rle(frame, mask_dir, categories)
assert new_frame.labels is not None
self.assertEqual(len(new_frame.labels), 11)
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 test_segtrack_to_rle(self) -> None:
"""Test seg_track to rle conversion."""
mask_dir = "./testcases/to_rle/seg_track/masks"
masks = ["0/0-1.jpg", "0/0-2.jpg"]
categories = get_leaf_categories(
load_bdd100k_config("seg_track").scalabel.categories
)
frames = [Frame(name=mask) for mask in masks]
new_frames = [
segtrack_to_rle(frame, mask_dir, categories) for frame in frames
]
for i in range(2):
labels = new_frames[i].labels
assert labels is not None
self.assertEqual(len(labels), 10)
self.assertEqual(new_frames[i].videoName, "0")
self.assertEqual(new_frames[i].frameIndex, i)
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 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)
