label outliers script now uses load data method defined in util

mu_map/random_search/label_outliers.py

@@ -10,6 +10,7 @@ from mu_map.dataset.transform import SequenceTransform, PadCropTranform
 from mu_map.models.unet import UNet
 from mu_map.random_search.cgan import load_params
 from mu_map.random_search.show_predictions import main
+from mu_map.random_search.eval.util import load_data
 
 controls = """
 Controls:
@@ -50,13 +51,7 @@ parser.add_argument(
 args = parser.parse_args()
 
 device = torch.device(args.device)
-
-runs = sorted(os.listdir(args.random_search_dir))
-runs = map(lambda f: os.path.join(args.random_search_dir, f), runs)
-runs = filter(lambda f: os.path.isdir(f), runs)
-runs = filter(lambda f: not os.path.islink(f), runs)
-runs = map(lambda f: os.path.basename(f), runs)
-runs = list(runs)
+runs = load_data(args.random_search_dir)
 
 data = {"run": [], "outlier": []}
 if os.path.isfile(args.out):
@@ -68,16 +63,17 @@ print(controls)
 
 total = str(len(runs))
 for i, run in enumerate(runs):
-    if int(run) in data["run"]:
+    print(f"Run {str(i+1):>{len(total)}}/{total}", end="\r")
+
+    if run in data["run"]:
         continue
 
-    print(f"Run {str(i+1):>{len(total)}}/{total}", end="\r")
     data["run"].append(int(run))
     data["outlier"].append(False)
 
-    _dir = os.path.join(args.random_search_dir, run)
 
-    params = load_params(os.path.join(_dir, "params.json"))
+    dir_run = os.path.join(args.random_search_dir, runs[run]["dir"])
+    params = runs[run]["params"]
 
     dataset = MuMapDataset(
         "data/second/",
@@ -90,7 +86,7 @@ for i, run in enumerate(runs):
 
     model = UNet(features=params["generator_features"])
     weights = torch.load(
-        os.path.join(_dir, "snapshots", "val_min_generator.pth"),
+        os.path.join(dir_run, "snapshots", "val_min_generator.pth"),
         map_location=device,
     )
     model.load_state_dict(weights)