add metadata to zarr attributes

scripts/inference.py

@@ -33,7 +33,11 @@ def _prepare_output_path(output_path: str | Path, force: bool) -> Path:
 
 def main():
     parser = argparse.ArgumentParser(description="Seismic feature inference using ViT foundation models.")
-    parser.add_argument("--model-path", required=True, help="Path to pretrained model directory (HuggingFace-style).")
+    parser.add_argument(
+        "--model-path",
+        required=True,
+        help="Path to pretrained model directory (HuggingFace-style). Such as NorskRegnesentralSTI/NCS-v1-2.5d-base.",
+    )
     parser.add_argument("--input-path", required=True, help="Path to input seismic file (.segy / .sgz).")
     parser.add_argument("--output-path", required=True, help="Output path (.zarr or .nc).")
     parser.add_argument("--direction", default="dir0", choices=["dir0", "dir90"], help="Primary traversal direction.")

src/NCS/inference/pipeline.py

@@ -4,6 +4,7 @@
 over tiled crops, and writes a patch-space feature cube via xarray + dask.
 """
 
+import importlib.metadata
 import logging
 import math
 import os
@@ -359,7 +360,11 @@ def run_inference(
     hidden_size = config.hidden_size
 
     logger.info(
-        "Loaded model type=%s (%s) hidden_size=%d from %s", model_type, str(type(model)), hidden_size, model_path
+        "Loaded model type=%s (%s) hidden_size=%d from %s",
+        base_model_type,
+        str(type(model)),
+        hidden_size,
+        model_path,
     )
 
     # 2. Open seismic file and compute stats
@@ -500,10 +505,39 @@ def run_inference(
 
     # 6. Save output
     output_path = Path(output_path)
+
+    already_written_directions = []
+    if output_path.exists():
+        with xr.open_dataset(output_path) as root_ds:
+            already_written_directions = root_ds.attrs.get("inference_directions", [])
+    if isinstance(already_written_directions, np.ndarray):
+        already_written_directions = already_written_directions.tolist()  # ty:ignore[no-matching-overload]
+    if isinstance(already_written_directions, str):
+        already_written_directions = [already_written_directions]
+    inference_direction = ":".join(input_views) if model_type == "vit25d" and input_views is not None else direction
+    inference_directions = sorted(set([inference_direction] + already_written_directions))
+
+    metadata = {
+        "ncs_version": importlib.metadata.version("NCS"),
+        "model_path": str(model_path),
+        "model_type": base_model_type,
+        "inference_directions": inference_directions,
+        "crop_size": [crop_size, crop_size, crop_size],
+        "extraction_mode": extraction_mode,
+        "num_overlap_patches": num_overlap_patches,
+        "inference_name": Path(str(model_path)).name,
+        "seismic_file": str(input_path),
+        "dtype": dtype,
+        "is_valid": 1,  # netCDF4 doesn't support boolean attributes, so we use an integer flag.
+    }
+
+    dt = xr.DataTree(name="root", children={direction: xr.DataTree(name=direction, dataset=ds)})
+    dt.attrs.update(metadata)
+
     if output_path.suffix == ".zarr":
-        ds.to_zarr(output_path, group=direction, mode="a")  # ty:ignore[invalid-argument-type]
+        dt.to_zarr(output_path, mode="a")
     else:
-        ds.to_netcdf(output_path, group=direction, mode="a")
+        dt.to_netcdf(output_path, mode="a")
 
     logger.info("Saved output to %s", output_path)
 

tests/test_pipeline.py

@@ -340,6 +340,82 @@ def barrier(self):
         assert captured["max_preload"] == 123
 
 
+@pytest.mark.parametrize("suffix", [".zarr", ".nc"])
+def test_run_inference_writes_fms_metadata_from_base_model_type_and_25d_views(tmp_path, monkeypatch, suffix):
+    if suffix == ".nc":
+        pytest.importorskip("netCDF4")
+    captured = {}
+
+    class _FakeModel:
+        def to(self, _device):
+            return self
+
+        def eval(self):
+            return self
+
+    def fake_runner(**kwargs):
+        captured["input_views"] = kwargs["input_views"]
+        return (
+            xr.Dataset(
+                {
+                    "features": (
+                        ["inline", "xline", "time_depth", "feature"],
+                        da.from_array(np.ones((1, 1, 1, 2), dtype=np.float32), chunks=(1, 1, 1, 2)),
+                    )
+                },
+                coords={"inline": [0], "xline": [0], "time_depth": [0], "feature": [0, 1]},
+            ),
+            None,
+        )
+
+    class _FakeSeismicFile:
+        def close(self):
+            return None
+
+    class _FakeWriter:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, tb):
+            return False
+
+        def barrier(self):
+            return None
+
+    monkeypatch.setattr(
+        pipeline.AutoConfig,
+        "from_pretrained",
+        lambda _path: SimpleNamespace(model_type="vit25d_mae", hidden_size=2, directions=["dir0", "dir45", "dir90"]),
+    )
+    monkeypatch.setattr(pipeline, "_load_model", lambda *args, **kwargs: _FakeModel())
+    monkeypatch.setattr(pipeline, "open_seismic", lambda _path: _FakeSeismicFile())
+    monkeypatch.setattr(pipeline, "cube_shape", lambda _f: (8, 8, 8))
+    monkeypatch.setattr(pipeline, "compute_cube_stats", lambda _f, n_traces=None: (0.0, 1.0))
+    monkeypatch.setattr(pipeline, "SeismicProcessor", lambda **kwargs: SimpleNamespace(**kwargs))
+    monkeypatch.setattr(pipeline, "_run_25d", fake_runner)
+    monkeypatch.setattr(pipeline, "AsyncWriter", _FakeWriter)
+
+    output_path = tmp_path / f"vit25d{suffix}"
+    pipeline.run_inference(
+        model_path="model",
+        input_path="cube.sgz",
+        output_path=output_path,
+        direction="dir0",
+        input_views=["dir0", "dir90"],
+    )
+
+    with xr.open_dataset(output_path) as root:
+        attrs = dict(root.attrs)
+
+    assert captured["input_views"] == ["dir0", "dir90"]
+    assert attrs["model_type"] == "vit25d_mae"
+    if suffix == ".zarr":
+        expected_directions = ["dir0:dir90"]
+    else:
+        expected_directions = "dir0:dir90"
+    assert attrs["inference_directions"] == expected_directions
+
+
 def test_run_inference_rejects_25d_input_views_without_primary_direction(tmp_path, monkeypatch):
     class _FakeModel:
         def to(self, _device):