diff --git a/datasets/cache_dataset3d.py b/datasets/cache_dataset3d.py index a06d55e..8e3e3c8 100644 --- a/datasets/cache_dataset3d.py +++ b/datasets/cache_dataset3d.py @@ -262,6 +262,8 @@ if __name__ == "__main__": from tutils.new.manager import ConfigManager config = ConfigManager() config.add_config("configs/vit_b.yaml") - dataset = Dataset3D(config=config['dataset'], use_cache=True) + + # Caching data + dataset = Dataset3D(config=config['dataset'], use_cache=False) dataset.caching_data() # dataset.convert_masks_types() diff --git a/datasets/cache_dataset3d3.py b/datasets/cache_dataset3d3.py index 2f5c972..f9d6bf3 100644 --- a/datasets/cache_dataset3d3.py +++ b/datasets/cache_dataset3d3.py @@ -33,7 +33,7 @@ class Dataset3D(basic_3d_dataset): if not os.path.isdir(dirpath): continue prefix = dirpath.split("/")[-1] - if prefix[:2] in config['cache_prefix']: + if prefix.split("_")[0] in config['cache_prefix']: data_paths += glob.glob(dirpath + "/label_jpg/*.jpg") print("Load ", dirpath) print('Masks len {}'.format(len(data_paths))) diff --git a/test/volume_eval.py b/test/volume_eval.py index c12fef4..02ef95e 100644 --- a/test/volume_eval.py +++ b/test/volume_eval.py @@ -141,7 +141,7 @@ if __name__ == "__main__": 'dataset_list': ['word'], # ["sabs"], chaos, word 'label_idx': 1, }, - 'pth': "model_latest.pth" + 'pth': "model.pth" } config = ConfigManager()