Add (final?) cvpr code + plotting scripts

Author: thetianshuhuang

dart/fields/grid.py

@@ -1,5 +1,7 @@
 """Basic Plenoxels-inspired Grid."""
 
+import os
+import numpy as np
 from jax import numpy as jnp
 import haiku as hk
 import math
@@ -73,10 +75,10 @@ def closure():
     def to_parser(p: types.ParserLike) -> None:
         """Create grid command line arguments."""
         p.add_argument(
-            "--lower", default=[-4.0, -4.0, -1.0], nargs='+', type=float,
+            "--lower", default=None, nargs='+', type=float,
             help="Lower coordinate (x, y, z).")
         p.add_argument(
-            "--upper", default=[4.0, 4.0, 1.0], nargs='+', type=float,
+            "--upper", default=None, nargs='+', type=float,
             help="Upper coordinate (x, y, z).")
         p.add_argument(
             "--resolution", default=25.0, type=float,
@@ -88,6 +90,12 @@ def to_parser(p: types.ParserLike) -> None:
     @classmethod
     def args_to_config(cls, args: types.Namespace) -> dict:
         """Create configuration dictionary."""
+        if args.lower is None or args.upper is None:
+            datadir = os.path.dirname(args.path)
+            npz = np.load(os.path.join(datadir, "map.npz"))
+            args.lower = npz['lower'].tolist()
+            args.upper = npz['upper'].tolist()
+
         assert len(args.upper) == 3
         assert len(args.lower) == 3
         grid_size = [

dart/metrics.py

@@ -19,11 +19,11 @@ def mse(
     Returns
     -------
     mse: MSE of the optimally-scaled `y_hat`.
-    alpha: Optimal scale factor.
+    xi: Optimal scale factor.
     """
-    alpha = jnp.sum(y_true * y_hat) / jnp.sum(y_hat**2)
-    mse = jnp.sum(jnp.square(y_true - alpha * y_hat))
-    return mse, alpha
+    xi = jnp.sum(y_true * y_hat) / jnp.sum(y_hat**2)
+    mse = jnp.sum(jnp.square(y_true - xi * y_hat))
+    return mse, xi
 
 
 def ssim(

notebooks/map.ipynb

@@ -0,0 +1,63 @@
+"""Result convenience wrapper."""
+
+import os
+import json
+import numpy as np
+import pandas as pd
+import h5py
+
+from beartype.typing import Any
+
+
+def _json(p):
+    with open(p) as f:
+        return json.load(f)
+
+
+class DartResult:
+    """DART result/dataset convenience wrapper."""
+
+    def __init__(self, path: str) -> None:
+        _meta = os.path.join(path, "metadata.json")
+        if not os.path.exists(_meta):
+            raise FileNotFoundError(
+                "Result path does not exist (could not find {})".format(_meta))
+
+        self.metadata = _json(_meta)
+        self.resdir = path
+        self.datadir = os.path.dirname(self.metadata["dataset"]["path"])
+
+    def dart(self) -> "DART":  # type: ignore
+        """Construct DART for results.
+        
+        NOTE: will import DART (and load jax & other heavy dependencies) on
+        first call.
+        """
+        from dart import DART
+        return DART.from_config(**self.metadata)
+
+    def path(self, subpath: str) -> str:
+        """Translate path to result/data directory."""
+        if subpath.startswith("data/"):
+            return os.path.join(self.datadir, subpath.replace("data/", ""))
+        else:
+            return os.path.join(self.resdir, subpath.replace("result/", ""))
+
+    def __getitem__(self, subpath: str) -> Any:
+        """Load npz/csv/h5/json file.
+        
+        Use `data/` to indicate when to load from the dataset directory;
+        otherwise, `subpath` is assumed to be in `results`. A `results/` prefix
+        can also be passed (which is removed).
+        """
+        path = self.path(subpath)
+        if not os.path.exists(path):
+            raise FileNotFoundError("File does not exist: {}".format(path))
+
+        def _err(p):
+            raise ValueError("Unknown file extension: {}".format(p))
+
+        exts = {
+            ".npz": np.load, ".csv": pd.read_csv,
+            ".h5": h5py.File, ".json": _json}
+        return exts.get(os.path.splitext(path)[1], _err)(path)

notebooks/volshow.ipynb

@@ -0,0 +1,37 @@
+"""Statistical utilities."""
+
+import os
+import numpy as np
+from jaxtyping import Num
+from beartype.typing import Union
+
+
+BASELINES = ['lidar', 'nearest', 'cfar', 'cfar_1e-2', 'cfar_1e-5', 'cfar_1e-8']
+
+
+def load_dir(path, key="ssim", baselines=BASELINES):
+    """Load metrics from a given dataset."""
+
+    def _try_load(*path):
+        try:
+            return np.load(os.path.join(*path))[key]
+        except FileNotFoundError:
+            return None
+
+    res = os.path.join("results", path)
+    data = os.path.join("data", path)
+
+    contents = {k: _try_load(res, k, "metrics.npz") for k in os.listdir(res)}
+    for k in baselines:
+        contents[k] = _try_load(data, "baselines/{}.npz".format(k))
+
+    return {k: v for k, v in contents.items() if v is not None}
+
+
+def effective_sample_size(x: Num[np.ndarray, "t"]) -> float:
+    """Calculate effective sample size for time series data."""
+    rho = np.array([
+        np.cov(x[i:], x[:-i])[0, 1] / np.std(x[i:]) / np.std(x[:-i])
+        for i in range(1, x.shape[0] // 2)])
+    rho_sum = np.sum(np.maximum(0.0, rho))
+    return x.shape[0] / (1 + 2 * rho_sum)

plot/_result.py

@@ -0,0 +1,38 @@
+from matplotlib import pyplot as plt
+import h5py
+import numpy as np
+
+
+def _im(ax, x):
+    lower, upper = np.percentile(x, [0, 99])
+    return ax.imshow(np.clip(x, lower, upper), cmap='inferno')
+
+h5 = h5py.File("results/boxes3/ngpsh/map.h5")
+fig, axs = plt.subplots(1, 2, figsize=(6, 4))
+im1 = _im(axs[0], np.mean(h5['sigma'][210:390, 230:410, 95:130], axis=2))
+im2 = _im(axs[1], np.mean(-h5['alpha'][210:390, 230:410, 95:130], axis=2))
+
+axs[0].text(5, 5, "Reflectance", color='white', ha='left', va='top')
+axs[1].text(5, 5, "Transmittance", color='white', ha='left', va='top')
+for ax in axs:
+    ax.set_xticks([])
+    ax.set_yticks([])
+    ax.text(50, 75, "(1)", color='white')
+    ax.text(125, 75, "(2)", color='white')
+    ax.text(20, 145, "(5)", color='white')
+    ax.text(78, 172, "(4)", color='white')
+    ax.text(135, 140, "(3)", color='white')
+fig.tight_layout(pad=1.0)
+
+cbar_ax = fig.add_axes([0.021, 0.08, 0.956, 0.04])
+fig.colorbar(im2, cax=cbar_ax, orientation='horizontal')
+cbar_ax.set_xticks([])
+cbar_ax.set_xlabel(r"Increasing Reflectance $\longrightarrow$", loc='left')
+
+cbar_ax2 = cbar_ax.twiny()
+cbar_ax2.xaxis.set_ticks_position('bottom')
+cbar_ax2.xaxis.set_label_position('bottom')
+cbar_ax2.set_xticks([])
+cbar_ax2.set_xlabel(r"$\longleftarrow$ Increasing Transmittance", loc='right')
+
+fig.savefig("figures/boxes.pdf", bbox_inches='tight')

plot/_stats.py

@@ -0,0 +1,73 @@
+"""Plot SSIM CDF."""
+
+import os
+from matplotlib import pyplot as plt
+import numpy as np
+
+from _stats import load_dir
+
+
+fig, axs = plt.subplots(2, 6, figsize=(12, 4))
+datasets = {
+    "boxes2": "Lab 1",
+    "boxes3": "Lab 2",
+    "wiselab4": "Office 1",
+    "wiselab5": "Office 2",
+    "mallesh-half": "Rowhouse 1",
+    "mallesh-1br": "Rowhouse 2",
+    "mallesh-full": "Rowhouse 3",
+    "agr-ground": "House 1",
+    "agr-full": "House 2",
+    "agr-yard": "Yard",
+    "tianshu-full": "Apartment 1",
+    "tianshu-half": "Apartment 2"
+}
+methods = {
+    "ngpsh": ("DART", 'C0', '-'),
+    "lidar": ("Lidar", 'C1', '--'),
+    "nearest": ("Nearest", 'C2', ':'),
+    "cfar": ("CFAR", 'C3', '-.')
+}
+
+for (ds, label), ax in zip(datasets.items(), axs.reshape(-1)):
+    ssim = load_dir(ds)
+    _ref = np.load(os.path.join("data", ds, "baselines", "reference.npz"))
+    ref = np.mean(_ref["ssim"], axis=0)
+
+    ax.axvline(
+        ref[0], color='black', linestyle='--', label='25/30/35db Reference',
+        linewidth=1.0)
+    ax.axvline(ref[1], color='black', linestyle='--', linewidth=1.0)
+    ax.axvline(ref[2], color='black', linestyle='--', linewidth=1.0)
+
+    for k, (desc, color, ls) in methods.items():
+        if k in ssim:
+            v = ssim[k]
+            ax.plot(
+                np.sort(v), np.arange(v.shape[0]) / v.shape[0],
+                label=desc, color=color, linestyle=ls)
+
+    ax.grid(visible=True)
+    ax.set_xlim(0.35, 0.85)
+    ax.text(0.83, 0.01, label, ha='right', va='bottom', backgroundcolor='white')
+
+for ax in axs[:, 1:].reshape(-1):
+    for tick in ax.yaxis.get_major_ticks():
+        tick.tick1line.set_visible(False)
+        tick.tick2line.set_visible(False)
+        tick.label1.set_visible(False)
+        tick.label2.set_visible(False)
+for ax in axs[:-1].reshape(-1):
+    for tick in ax.xaxis.get_major_ticks():
+        tick.tick1line.set_visible(False)
+        tick.tick2line.set_visible(False)
+        tick.label1.set_visible(False)
+        tick.label2.set_visible(False)
+
+
+fig.tight_layout(h_pad=0.2, w_pad=0.5)
+axs[-1, -1].legend(
+    ncols=5, loc='upper right', bbox_to_anchor=(1.05, -0.15), frameon=False)
+axs[1, 0].set_ylabel(r"Cumulative Probability $\longrightarrow$", loc='bottom')
+axs[1, 0].set_xlabel(r"SSIM (higher is better) $\longrightarrow$", loc='left')
+fig.savefig("figures/ssim.pdf", bbox_inches='tight')