fix issues in predict and train

aramis-lab · Jun 4, 2024 · e8ddf96 · e8ddf96
1 parent 4efbcac
commit e8ddf96
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/tests/test_predict.py b/tests/test_predict.py
@@ -117,6 +117,6 @@ def test_predict(cmdopt, tmp_path, test_name):
 
     assert compare_folders(
         tmp_out_dir / maps_name,
-        ref_dir / maps_name,
+        input_dir / maps_name,
         tmp_out_dir,
     )
diff --git a/tests/test_train_ae.py b/tests/test_train_ae.py
@@ -33,8 +33,10 @@ def test_train_ae(cmdopt, tmp_path, test_name):
 
     labels_path = str(input_dir / "labels_list" / "2_fold")
     config_path = str(input_dir / "train_config.toml")
+    split = 0
+
     if test_name == "image_ae":
-        split = [0, 0]
+        split = 1
         test_input = [
             "train",
             "reconstruction",
@@ -45,10 +47,9 @@ def test_train_ae(cmdopt, tmp_path, test_name):
             "-c",
             config_path,
             "--split",
-            "1",
+            str(split),
         ]
     elif test_name == "patch_multi_ae":
-        split = [0, 0]
         test_input = [
             "train",
             "reconstruction",
@@ -61,7 +62,6 @@ def test_train_ae(cmdopt, tmp_path, test_name):
             "--multi_network",
         ]
     elif test_name == "roi_ae":
-        split = [0, 0]
         test_input = [
             "train",
             "reconstruction",
@@ -73,7 +73,6 @@ def test_train_ae(cmdopt, tmp_path, test_name):
             config_path,
         ]
     elif test_name == "slice_ae":
-        split = [0, 0]
         test_input = [
             "train",
             "reconstruction",
@@ -116,7 +115,7 @@ def test_train_ae(cmdopt, tmp_path, test_name):
         tmp_path,
     )
     assert compare_folders(
-        tmp_out_dir / f"split-{split[0]}" / "best-loss",
-        ref_dir / ("maps_" + test_name) / f"split-{split[1]}" / "best-loss",
+        tmp_out_dir / f"split-{split}" / "best-loss",
+        ref_dir / ("maps_" + test_name) / f"split-{split}" / "best-loss",
         tmp_path,
     )