deform_asset.py

"""
Deform an implicit asset with feature-field handles, bake the result into a new
asset, and render it.

    python deform_asset.py <asset> <encoder.pth> <handles.json> <outdir>

``<asset>`` is a 3DGS `.ply`, a point-cloud `.ply` (auto-detected), or a
nerfstudio `config.yml` (needs nerfstudio). ``<encoder.pth>`` is the feature
field distilled *into that asset* by the matching ``distillation_*.py`` script.
``<handles.json>`` lists exact 3D handle points with full affine transforms —
see ``deformfield.py`` for the format, or ``assets/spot/handles_example.json``.

What gets baked per representation:

  point cloud   points forward-warped              -> {outdir}/deformed.ply
  gaussians     centers forward-warped; covariances
                transformed by the warp's local
                Jacobian (Sigma' = J Sigma J^T)     -> {outdir}/deformed.ply
  NeRF          nothing to move — volumetric sample
                queries are backward-warped instead -> renders only (a NeRF
                cannot be re-baked without retraining)

All inputs are assumed to already be in the cube-normalized frame the encoder
was distilled in (the ``distillation_*.py`` scripts print/save that transform).
Rendered views of the rest and deformed asset land in {outdir}/rest/ and
{outdir}/deformed/.
"""
import os
import sys
import argparse
import numpy as np
import torch

import gsply
import deformfield
import repdistill
from distillation_pc import rasterize_index, _focal

_HERE = os.path.dirname(os.path.abspath(__file__))


def load_encoder(pth, device):
    sys.path.insert(0, os.path.join(_HERE, "scripts", "reptransfer"))
    from encoder import load_encoder as _load
    return _load(pth, device)


def is_gaussian_ply(path):
    """3DGS plys carry scale_/rot_ vertex properties; plain clouds don't."""
    from plyfile import PlyData
    return "scale_0" in PlyData.read(path)["vertex"].data.dtype.names


def render_gs(g, w2c_cv, K, H, W, outdir, device="cuda"):
    imgs = []
    for i in range(len(w2c_cv)):
        rgb, _ = gsply.rasterize(g, w2c_cv[i:i + 1], K[i:i + 1], H, W, device=device)
        imgs.append(rgb[0].clamp(0, 1).cpu().numpy())
    repdistill_save(imgs, outdir)


def render_pc(points, colors, w2c, fl, cx, cy, H, W, outdir, radius=1):
    imgs = []
    for i in range(len(w2c)):
        idx = rasterize_index(points, w2c[i], fl, fl, cx, cy, H, W, radius)
        img = np.ones((H, W, 3))
        hit = idx >= 0
        img[hit] = colors[idx[hit]]
        imgs.append(img)
    repdistill_save(imgs, outdir)


def repdistill_save(imgs, outdir):
    from PIL import Image
    os.makedirs(outdir, exist_ok=True)
    for i, img in enumerate(imgs):
        Image.fromarray((np.clip(img, 0, 1) * 255).astype(np.uint8)).save(
            os.path.join(outdir, f"view_{i:02d}.png"))
    print(f"wrote {len(imgs)} views -> {outdir}")


def report_weights(encoder, hpoints, feats, sample_points):
    """Print each handle's weight stats so a dead handle is obvious immediately."""
    for i, (hp, f_h) in enumerate(zip(hpoints, feats)):
        w = deformfield.weight_field(encoder, sample_points, f_h)
        print(f"handle {i} @ {hp.tolist()}  w: max={w.max():.3f} "
              f"mean={w.mean():.3f} frac>0.5={(w > 0.5).float().mean():.3f}")
        if float(w.max()) < 0.1:
            print(f"  WARNING: handle {i} moves almost nothing — is the point "
                  f"on the object, in the cube-normalized frame?")


def deform_gaussians(g, encoder, feats, affines):
    """Warp centers and covariances. Returns a new gaussian dict."""
    means = g["means"]
    means_def = deformfield.forward_warp(encoder, means, feats, affines)

    out = dict(g)
    out["means"] = means_def.cpu().float()

    J = deformfield.local_jacobian(encoder, means, feats, affines)      # (N,3,3)
    if not torch.allclose(J, torch.eye(3).expand_as(J)):
        scales = torch.exp(g["scales"])
        R = gsply.quat_to_rotmat(g["quats"])
        Sigma = (R * (scales ** 2)[:, None, :]) @ R.transpose(1, 2)     # R diag(s^2) R^T
        Sigma = J @ Sigma @ J.transpose(1, 2)                           # J Sigma J^T
        evals, evecs = torch.linalg.eigh(Sigma)
        det = torch.linalg.det(evecs)                                   # keep rotations proper
        evecs = torch.where((det < 0)[:, None, None],
                            evecs * torch.tensor([-1.0, 1.0, 1.0]), evecs)
        out["scales"] = torch.log(torch.sqrt(torch.clamp(evals, min=1e-12))).float()
        out["quats"] = gsply.rotmat_to_quat(evecs).float()
        print("transformed covariances by the warp's local Jacobian")
    return out


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("asset", type=str, help="3DGS .ply, point-cloud .ply, or nerfstudio config.yml")
    ap.add_argument("encoder", type=str, help="final_encoder.pth distilled into this asset")
    ap.add_argument("handles", type=str, help="handles JSON (see deformfield.py)")
    ap.add_argument("outdir", type=str)
    ap.add_argument("--nviews", type=int, default=12)
    ap.add_argument("--viewtype", type=str, choices=["default", "fib", "ring"], default="fib",
                    help="ring = turntable at 30 deg elevation")
    ap.add_argument("--viewradius", type=float, default=2.8)
    ap.add_argument("--res", type=int, default=512)
    ap.add_argument("--radius", type=int, default=1, help="point-cloud splat radius (pixels)")
    ap.add_argument("--iters", type=int, default=10,
                    help="NeRF inverse-warp fixed-point iteration cap (early-stops on convergence)")
    ap.add_argument("--warp_grid", type=int, default=192,
                    help="NeRF only: bake the inverse warp on this-resolution voxel grid "
                         "for fast lookup; 0 = exact per-query warping (slow)")
    ap.add_argument("--no_rest", action="store_true", help="skip rendering the rest state")
    args = ap.parse_args()

    device = "cuda"
    os.makedirs(args.outdir, exist_ok=True)
    H = W = args.res

    encoder = load_encoder(args.encoder, device)
    hpoints, affines = deformfield.load_handles(args.handles)
    feats = deformfield.handle_features(encoder, hpoints)
    print(f"{len(feats)} handle(s) loaded from {args.handles}")

    eyes, w2c_gl = repdistill.cameras(args.nviews, args.viewtype, args.viewradius)
    fl, _, cx, cy = _focal(H, W)

    if args.asset.endswith((".yml", ".yaml")):
        # ---- NeRF: backward-warp every volumetric sample query ----
        sys.path.insert(0, os.path.join(_HERE, "scripts", "reptransfer"))
        from nsload import load_model_only, eval_camera_batch, render_view, over_white
        from nerfstudio.cameras.rays import Frustums
        from nerfstudio.cameras.cameras import Cameras, CameraType

        model = load_model_only(args.asset, device=device)
        c2w = np.stack([np.linalg.inv(m) for m in w2c_gl])
        cameras = Cameras(
            camera_to_worlds=torch.from_numpy(c2w[:, :3, :]).float(),
            fx=fl, fy=fl, cx=cx, cy=cy, width=W, height=H,
            camera_type=CameraType.PERSPECTIVE,
        ).to(device)

        def render_all(outdir):
            import time
            imgs = []
            for i in range(len(c2w)):
                t0 = time.time()
                rgb, _, acc = render_view(model, cameras, i)
                imgs.append(over_white(rgb, acc))
                print(f"  view {i + 1}/{len(c2w)} in {time.time() - t0:.1f}s", flush=True)
            repdistill_save(imgs, outdir)

        report_weights(encoder, hpoints, feats,
                       torch.rand(20000, 3) * 2 - 1)   # probe the cube volume

        if not args.no_rest:
            render_all(os.path.join(args.outdir, "rest"))

        # A view warps ~10^8 sample positions; bake the (smooth) inverse warp on a
        # grid once and answer every query with a trilinear lookup instead of an
        # encoder evaluation. --warp_grid 0 falls back to exact per-query warping.
        if args.warp_grid > 0:
            warp = deformfield.bake_backward_grid(
                encoder, feats, affines, res=args.warp_grid, iters=args.iters)
            probe = torch.rand(20000, 3, device=device) * 2 - 1
            exact = deformfield.backward_warp(encoder, probe, feats, affines,
                                              iters=args.iters)
            err = (warp(probe) - exact).norm(dim=1).max()
            print(f"grid-vs-exact max error: {float(err):.5f} (cube units)")
        else:
            def warp(y):
                return deformfield.backward_warp(encoder, y, feats, affines,
                                                 iters=args.iters)

        _orig = Frustums.get_positions

        def warped_get_positions(self):
            pos = _orig(self)
            return warp(pos.reshape(-1, 3)).reshape(pos.shape).to(pos.dtype)

        Frustums.get_positions = warped_get_positions
        try:
            render_all(os.path.join(args.outdir, "deformed"))
        finally:
            Frustums.get_positions = _orig
        print("NeRF cannot be re-baked without retraining; deformed renders are the output.")
        return

    # ---- .ply: gaussian splat or point cloud ----
    K = gsply.intrinsics_matrix(fl, fl, cx, cy, n=len(w2c_gl))
    w2c_cv = gsply.gl_to_cv(w2c_gl)
    outply = os.path.join(args.outdir, "deformed.ply")

    if is_gaussian_ply(args.asset):
        g = gsply.load_ply(args.asset)
        report_weights(encoder, hpoints, feats, g["means"])
        if not args.no_rest:
            render_gs(g, w2c_cv, K, H, W, os.path.join(args.outdir, "rest"), device)
        g_def = deform_gaussians(g, encoder, feats, affines)
        gsply.save_ply(outply, g_def)
        render_gs(g_def, w2c_cv, K, H, W, os.path.join(args.outdir, "deformed"), device)
    else:
        import trimesh
        pc = trimesh.load(args.asset, process=False)
        points = np.asarray(pc.vertices, dtype=np.float64)
        try:
            colors = np.asarray(pc.colors)[:, :3].astype(np.float64) / 255.0
            assert colors.shape[0] == points.shape[0]
        except Exception:
            colors = np.full_like(points, 0.7)
        report_weights(encoder, hpoints, feats, torch.from_numpy(points).float())
        if not args.no_rest:
            render_pc(points, colors, w2c_gl, fl, cx, cy, H, W,
                      os.path.join(args.outdir, "rest"), args.radius)
        pts_def = deformfield.forward_warp(encoder, points, feats, affines).cpu().numpy()
        trimesh.PointCloud(vertices=pts_def.astype(np.float32),
                           colors=(colors * 255).astype(np.uint8)).export(outply)
        print(f"wrote {len(pts_def)} points -> {outply}")
        render_pc(pts_def, colors, w2c_gl, fl, cx, cy, H, W,
                  os.path.join(args.outdir, "deformed"), args.radius)


if __name__ == "__main__":
    main()