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)