def create_fake_detection_data_loader()

in D2Go/create_d2go.py [0:0]

• 24 lines of code
• 2 McCabe index (conditional complexity)

def create_fake_detection_data_loader(height, width, is_train):
    with make_temp_directory("detectron2go_tmp_dataset") as dataset_dir:
        runner = create_runner("d2go.runner.GeneralizedRCNNRunner")
        cfg = runner.get_default_cfg()
        cfg.DATASETS.TRAIN = ["default_dataset_train"]
        cfg.DATASETS.TEST = ["default_dataset_test"]

        with make_temp_directory("detectron2go_tmp_dataset") as dataset_dir:
            image_dir = os.path.join(dataset_dir, "images")
            os.makedirs(image_dir)
            image_generator = LocalImageGenerator(image_dir, width=width, height=height)

            if is_train:
                with register_toy_dataset(
                    "default_dataset_train", image_generator, num_images=3
                ):
                    train_loader = runner.build_detection_train_loader(cfg)
                    yield train_loader
            else:
                with register_toy_dataset(
                    "default_dataset_test", image_generator, num_images=3
                ):
                    test_loader = runner.build_detection_test_loader(
                        cfg, dataset_name="default_dataset_test"
                    )
                    yield test_loader