Make full evaluation optional in infer script

Add an --eval flag to infer.py and default to single-sample inference to keep Colab runs fast.

README.md

@@ -69,6 +69,12 @@ python infer.py CLDM  # SAII-CLDM
 `infer.py` uses the bundled Overthrust sample and writes outputs under
 `outputs/infer_LDDPM/` or `outputs/infer_CLDM/`.
 
+Add `--eval` to run the full bundled Overthrust evaluation:
+
+```bash
+python infer.py CLDM --eval
+```
+
 ## Overthrust Results
 
 Impedance-domain metrics on the bundled Overthrust setting:

infer.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import argparse
 import sys
 from pathlib import Path
 
@@ -16,11 +17,15 @@ from codes.pipeline import SeismicImpInvCLDMPipeline, SeismicImpInvLDDPMPipeline
 from codes.util import OverthrustForwardOperator, ricker_wavelet
 
 
-METHOD = sys.argv[1].upper() if len(sys.argv) > 1 else "LDDPM"
-OUT_DIR = REPO_ROOT / "outputs" / f"infer_{METHOD}"
 PATCH_INDEX = 0
 MODEL_DIR = REPO_ROOT
-RUN_EVAL = True
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run SAII-LDDPM/CLDM inference.")
+    parser.add_argument("method", nargs="?", choices=["LDDPM", "CLDM"], default="LDDPM")
+    parser.add_argument("--eval", action="store_true", help="Run full Overthrust evaluation after single-sample inference.")
+    return parser.parse_args()
 
 def save_comparison(dipin, record, target, prediction, output_path):
     fig, axes = plt.subplots(1, 4, figsize=(16, 4))
@@ -40,18 +45,19 @@ def save_comparison(dipin, record, target, prediction, output_path):
 
 
 if __name__ == "__main__":
-    if METHOD not in {"LDDPM", "CLDM"}:
-        raise ValueError("METHOD must be LDDPM or CLDM. Example: python infer.py CLDM")
+    args = parse_args()
+    method = args.method.upper()
+    out_dir = REPO_ROOT / "outputs" / f"infer_{method}"
 
-    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    out_dir.mkdir(parents=True, exist_ok=True)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Using device: {device}")
-    print(f"Method: {METHOD}")
+    print(f"Method: {method}")
 
     dataset = OverthrustTrueimpDataset(
         patch_indices=[PATCH_INDEX],
         data_dir=REPO_ROOT / "data",
-        cache_dir=OUT_DIR / "cache",
+        cache_dir=out_dir / "cache",
     )
     sample = dataset[0]
     dipin = sample["dipin"].unsqueeze(0).to(device)
@@ -59,7 +65,7 @@ if __name__ == "__main__":
     image = sample["image"].unsqueeze(0).to(device)
     seed = int(sample["seed"])
 
-    if METHOD == "LDDPM":
+    if method == "LDDPM":
         num_inference_steps = 1000
         extra_kwargs = {}
         pipe = SeismicImpInvLDDPMPipeline.from_pretrained(
@@ -109,15 +115,15 @@ if __name__ == "__main__":
     dipin_np = dipin[0, 0].detach().cpu().numpy()
     record_np = record[0, 0].detach().cpu().numpy()
 
-    np.save(OUT_DIR / "prediction.npy", prediction)
-    np.save(OUT_DIR / "target.npy", target)
-    save_comparison(dipin_np, record_np, target, prediction, OUT_DIR / "comparison.png")
+    np.save(out_dir / "prediction.npy", prediction)
+    np.save(out_dir / "target.npy", target)
+    save_comparison(dipin_np, record_np, target, prediction, out_dir / "comparison.png")
 
-    print(f"Saved: {OUT_DIR / 'prediction.npy'}")
-    print(f"Saved: {OUT_DIR / 'target.npy'}")
-    print(f"Saved: {OUT_DIR / 'comparison.png'}")
+    print(f"Saved: {out_dir / 'prediction.npy'}")
+    print(f"Saved: {out_dir / 'target.npy'}")
+    print(f"Saved: {out_dir / 'comparison.png'}")
 
-    if RUN_EVAL:
+    if args.eval:
         from codes.eval_overthrust import evaluate_overthrust
 
-        evaluate_overthrust(pipe, method=METHOD, output_dir=OUT_DIR / "eval")
+        evaluate_overthrust(pipe, method=method, output_dir=out_dir / "eval")