| """General Lightning-based inference script for TactileVAE. |
| |
| Features: |
| - Load any Lightning `.ckpt` checkpoint. |
| - Load any config YAML. |
| - Randomly select `N` samples from any split (`train` / `val` / `test`). |
| - Run reconstruction inference and save metrics + visualization. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import pytorch_lightning as pl |
| import torch |
| import yaml |
| from PIL import Image |
| from torch.utils.data import DataLoader, Subset |
|
|
| _REPO_ROOT = Path(__file__).resolve().parents[2] |
| if str(_REPO_ROOT) not in sys.path: |
| sys.path.insert(0, str(_REPO_ROOT)) |
|
|
| from tactile_vae.dataset import TactileParquetDataset |
| from tactile_vae.model import TactileVAE |
|
|
| DEFAULT_CONFIG = Path("/group2/ct/weihanx/tactile_world_model/runs/vae_baseline_3/config.snapshot.yaml") |
| DEFAULT_CKPT = Path("/group2/ct/weihanx/tactile_world_model/runs/vae_baseline_3/checkpoints/last.ckpt") |
| DEFAULT_OUT_DIR = Path("/group2/ct/weihanx/tactile_world_model/tactile_vae/inference/vae_baseline_3") |
|
|
|
|
| def _resolve_path(p: str | Path) -> Path: |
| path = Path(p) |
| return path if path.is_absolute() else (_REPO_ROOT / path).resolve() |
|
|
|
|
| def load_config(path: Path) -> dict: |
| with path.open() as f: |
| cfg = yaml.safe_load(f) |
| if not isinstance(cfg, dict): |
| raise ValueError(f"invalid config: {path}") |
| cfg["data"]["root"] = str(_resolve_path(cfg["data"]["root"])) |
| if cfg["data"].get("splits_path"): |
| cfg["data"]["splits_path"] = str(_resolve_path(cfg["data"]["splits_path"])) |
| return cfg |
|
|
|
|
| def pick_device(spec: str) -> torch.device: |
| if spec == "auto": |
| return torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| return torch.device(spec) |
|
|
|
|
| class InferenceModule(pl.LightningModule): |
| """Minimal LightningModule used for strict Lightning checkpoint loading.""" |
|
|
| def __init__(self, config: dict): |
| super().__init__() |
| self.config = config |
| self.model = TactileVAE(**config["model"]) |
|
|
| def forward(self, x, **kw): |
| return self.model(x, **kw) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| p = argparse.ArgumentParser() |
| p.add_argument("--config", type=Path, default=DEFAULT_CONFIG, help="config yaml") |
| p.add_argument("--ckpt", type=Path, default=DEFAULT_CKPT, help="Lightning checkpoint .ckpt") |
| p.add_argument("--out-dir", type=Path, default=DEFAULT_OUT_DIR, help="output directory") |
| p.add_argument("--split", type=str, default="test", choices=["train", "val", "test"]) |
| p.add_argument("--num-samples", type=int, default=50, help="number of random samples from the split") |
| p.add_argument("--batch-size", type=int, default=16) |
| p.add_argument("--num-workers", type=int, default=0) |
| p.add_argument("--seed", type=int, default=0) |
| p.add_argument("--device", type=str, default="auto", help="auto / cuda / cpu / cuda:0 ...") |
| p.add_argument("--max-grid", type=int, default=16, help="max samples shown in saved reconstruction grid") |
| return p.parse_args() |
|
|
|
|
| def build_dataset(cfg: dict, split: str) -> TactileParquetDataset: |
| dcfg = cfg["data"] |
| return TactileParquetDataset( |
| root=dcfg["root"], |
| split=split, |
| splits_path=dcfg.get("splits_path"), |
| image_size=dcfg["image_size"], |
| cache_files=dcfg.get("cache_files", 1), |
| color_jitter=None, |
| ) |
|
|
|
|
| def select_subset(ds: TactileParquetDataset, n: int, seed: int) -> tuple[Subset, list[int]]: |
| n = min(max(1, int(n)), len(ds)) |
| rng = np.random.default_rng(seed) |
| idx = rng.choice(len(ds), size=n, replace=False).tolist() |
| return Subset(ds, idx), idx |
|
|
|
|
| @torch.no_grad() |
| def run_inference( |
| module: InferenceModule, |
| ds: TactileParquetDataset, |
| subset_idx: list[int], |
| loader: DataLoader, |
| device: torch.device, |
| ) -> tuple[list[dict[str, Any]], float, float, list[tuple[torch.Tensor, torch.Tensor]]]: |
| module.eval().to(device) |
| per_sample: list[dict[str, Any]] = [] |
| vis_pairs: list[tuple[torch.Tensor, torch.Tensor]] = [] |
| mae_total = 0.0 |
| mse_total = 0.0 |
| n_total = 0 |
|
|
| cursor = 0 |
| for x in loader: |
| x = x.to(device, non_blocking=True) |
| out = module.model(x, sample=False) |
| x_hat = out["x_hat"] |
|
|
| abs_err = (x - x_hat).abs().mean(dim=(1, 2, 3)) |
| sq_err = ((x - x_hat) ** 2).mean(dim=(1, 2, 3)) |
| bs = x.shape[0] |
|
|
| for i in range(bs): |
| gidx = subset_idx[cursor + i] |
| sample_id = ds.sample_id(gidx) |
| mae_i = float(abs_err[i].item()) |
| mse_i = float(sq_err[i].item()) |
| per_sample.append( |
| { |
| "subset_rank": cursor + i, |
| "dataset_index": int(gidx), |
| "sample_id": sample_id, |
| "mae": mae_i, |
| "mse": mse_i, |
| } |
| ) |
| vis_pairs.append((x[i].detach().cpu(), x_hat[i].detach().cpu())) |
| mae_total += mae_i |
| mse_total += mse_i |
| n_total += 1 |
| cursor += bs |
|
|
| mae_mean = mae_total / max(1, n_total) |
| mse_mean = mse_total / max(1, n_total) |
| return per_sample, mae_mean, mse_mean, vis_pairs |
|
|
|
|
| def save_grid(pairs: list[tuple[torch.Tensor, torch.Tensor]], out_path: Path, n_show: int, image_size: int) -> None: |
| n = min(n_show, len(pairs)) |
| if n <= 0: |
| return |
| h = w = int(image_size) |
| canvas = np.zeros((2 * h, n * w, 3), dtype=np.uint8) |
| for i in range(n): |
| src, rec = pairs[i] |
| src_np = (src.clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8) |
| rec_np = (rec.clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype(np.uint8) |
| canvas[:h, i * w : (i + 1) * w] = src_np |
| canvas[h:, i * w : (i + 1) * w] = rec_np |
| out_path.parent.mkdir(parents=True, exist_ok=True) |
| Image.fromarray(canvas).save(out_path) |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| cfg = load_config(args.config) |
| device = pick_device(args.device) |
|
|
| args.out_dir.mkdir(parents=True, exist_ok=True) |
| print(f"config: {args.config}") |
| print(f"ckpt: {args.ckpt}") |
| print(f"split: {args.split}") |
| print(f"num_samples: {args.num_samples}") |
| print(f"device: {device}") |
| print(f"out_dir: {args.out_dir}") |
|
|
| ds = build_dataset(cfg, split=args.split) |
| subset, subset_idx = select_subset(ds, args.num_samples, args.seed) |
| print(f"split_size={len(ds)} selected={len(subset_idx)}") |
| print(f"preview_sample_ids={[ds.sample_id(i) for i in subset_idx[:5]]}") |
|
|
| loader = DataLoader( |
| subset, |
| batch_size=min(max(1, args.batch_size), len(subset)), |
| shuffle=False, |
| num_workers=args.num_workers, |
| pin_memory=device.type == "cuda", |
| drop_last=False, |
| persistent_workers=args.num_workers > 0, |
| ) |
|
|
| module = InferenceModule.load_from_checkpoint( |
| str(args.ckpt), |
| config=cfg, |
| strict=True, |
| map_location="cpu", |
| ) |
|
|
| per_sample, mae_mean, mse_mean, vis_pairs = run_inference( |
| module=module, ds=ds, subset_idx=subset_idx, loader=loader, device=device |
| ) |
|
|
| grid_path = args.out_dir / "reconstruction_grid.png" |
| save_grid(vis_pairs, out_path=grid_path, n_show=args.max_grid, image_size=cfg["data"]["image_size"]) |
|
|
| summary = { |
| "config": str(args.config), |
| "checkpoint": str(args.ckpt), |
| "split": args.split, |
| "seed": args.seed, |
| "selected_num_samples": len(subset_idx), |
| "mean_mae": mae_mean, |
| "mean_mse": mse_mean, |
| "grid_path": str(grid_path), |
| } |
| with (args.out_dir / "summary.json").open("w") as f: |
| json.dump(summary, f, indent=2) |
| with (args.out_dir / "per_sample_metrics.json").open("w") as f: |
| json.dump(per_sample, f, indent=2) |
|
|
| print(f"mean_mae={mae_mean:.6f} mean_mse={mse_mean:.6f}") |
| print(f"saved: {args.out_dir / 'summary.json'}") |
| print(f"saved: {args.out_dir / 'per_sample_metrics.json'}") |
| print(f"saved: {grid_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
