From ce31f13be07535bf7e3d2d6f270c9f68c5eb8cab Mon Sep 17 00:00:00 2001
From: transcendentsky <transcendentsky@gmail.com>
Date: Wed, 20 Mar 2024 22:07:35 +0800
Subject: [PATCH] .

---
 datasets/cache_dataset3d.py  | 4 +++-
 datasets/cache_dataset3d3.py | 2 +-
 test/volume_eval.py          | 2 +-
 3 files changed, 5 insertions(+), 3 deletions(-)

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()