def _create_dataset()

in activemri/envs/envs.py [0:0]

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

    def _create_dataset(self) -> DataInitFnReturnType:
        root_path = pathlib.Path(self._data_location)
        datacache_dir = activemri.envs.util.maybe_create_datacache_dir()

        train_path = root_path / f"{self.dataset_name}_train"
        val_path = root_path / f"{self.dataset_name}_val"
        val_cache_file = datacache_dir / f"val_{self.dataset_name}_cache.pkl"
        test_path = root_path / f"{self.dataset_name}_test"
        test_cache_file = datacache_dir / f"test_{self.dataset_name}_cache.pkl"

        if not test_path.is_dir():
            warnings.warn(
                f"No test directory found for {self.dataset_name}. "
                f"I will use val directory for test model (env.set_test())."
            )
            test_path = val_path
            test_cache_file = val_cache_file

        train_data = fastmri.data.SliceDataset(
            train_path,
            ActiveMRIEnv._void_transform,
            challenge=self.challenge,
            num_cols=self.num_cols,
            dataset_cache_file=datacache_dir / f"train_{self.dataset_name}_cache.pkl",
        )
        val_data = fastmri.data.SliceDataset(
            val_path,
            ActiveMRIEnv._void_transform,
            challenge=self.challenge,
            num_cols=self.num_cols,
            dataset_cache_file=val_cache_file,
        )
        test_data = fastmri.data.SliceDataset(
            test_path,
            ActiveMRIEnv._void_transform,
            challenge=self.challenge,
            num_cols=self.num_cols,
            dataset_cache_file=test_cache_file,
        )
        return train_data, val_data, test_data