make code more flexible

Author: SarahMuth

run_sbatch_revision.sbatch

@@ -2,13 +2,13 @@
 #SBATCH -c 4 #4 #8
 #SBATCH --mem 120G #120G #32G #64G #256G
 #SBATCH -p grete:shared #grete:shared #grete-h100:shared
-#SBATCH -t 3:00:00 #6:00:00 #48:00:00 
+#SBATCH -t 2:00:00 #6:00:00 #48:00:00 
 #SBATCH -G A100:1 #V100:1 #2 #A100:1  #gtx1080:2 #v100:1 #H100:1
 #SBATCH --output=/user/muth9/u12095/synapse-net/slurm_revision/slurm-%j.out 
 #SBATCH -A nim00007 #SBATCH --constraint 80gb 
 
 source ~/.bashrc
 conda activate synapse-net
-python /user/muth9/u12095/synapse-net/scripts/cooper/revision/thin_az_gt.py \
- -i /mnt/ceph-hdd/cold/nim00007/new_AZ_train_data/stem_for_eval/ \
- -o /mnt/ceph-hdd/cold/nim00007/new_AZ_train_data/stem_for_eval/
+python /user/muth9/u12095/synapse-net/scripts/cooper/revision/merge_az.py -v 6 \
+ --names stem_for_eval \
+ -i /mnt/ceph-hdd/cold/nim00007/new_AZ_train_data/

scripts/cooper/revision/az_prediction.py

@@ -54,10 +54,12 @@ def run_prediction(model, name, split_folder, version, split_names, in_path):
 
 
 def get_model(version):
-    assert version in (3, 4, 5, 7)
+    assert version in (3, 4, 5, 6, 7)
     split_folder = get_split_folder(version)
     if version == 3:
         model_path = os.path.join(split_folder, "checkpoints", "3D-AZ-model-TEM_STEM_ChemFix_wichmann-v3")
+    elif version ==6:
+        model_path = "/mnt/ceph-hdd/cold/nim00007/models/AZ/v6/"
     elif version == 7:
         model_path = "/mnt/lustre-emmy-hdd/usr/u12095/synapse_net/models/ConstantinAZ/checkpoints/v7/"
     else:

scripts/cooper/revision/merge_az.py

@@ -18,14 +18,18 @@
 
 
 # STEM CROPPED IS OFTEN TOO SMALL!
-def merge_az(name, version, check):
+def merge_az(name, version, check, in_path):
     split_folder = get_split_folder(version)
 
     if name == "stem_cropped":
         file_paths = glob(os.path.join("/mnt/ceph-hdd/cold/nim00007/new_AZ_train_data/stem_cropped", "*.h5"))
         file_names = [os.path.basename(path) for path in file_paths]
     else:
-        file_names = get_file_names(name, split_folder, split_names=["train", "val", "test"])
+        if in_path:
+            file_paths = glob(os.path.join(in_path, name, "*.h5"))
+            file_names = [os.path.basename(path) for path in file_paths]
+        else:
+            file_names = get_file_names(name, split_folder, split_names=["train", "val", "test"])
     seg_paths, gt_paths = get_paths(name, file_names)
 
     for seg_path, gt_path in zip(seg_paths, gt_paths):
@@ -79,7 +83,7 @@ def visualize_merge(args):
     for name in args.names:
         if "endbulb" in name:
             continue
-        merge_az(name, args.version, check=True)
+        merge_az(name, args.version, check=True, in_path=args.in_path)
 
 
 def copy_az(name, version):
@@ -101,7 +105,7 @@ def run_merge(args):
         if "endbulb" in name:
             copy_az(name, args.version)
         else:
-            merge_az(name, args.version, check=False)
+            merge_az(name, args.version, check=False, in_path= args.in_path)
 
 
 def main():
@@ -110,6 +114,7 @@ def main():
     parser.add_argument("--visualize", action="store_true")
     parser.add_argument("--names", nargs="+", default=ALL_NAMES + ["endbulb_of_held_cropped"])
     parser.add_argument("--version", "-v", type=int, default=4)
+    parser.add_argument("--in_path", "-i", default=None)
 
     args = parser.parse_args()
     if args.visualize: