memef4rmer/edit_anything

1import numpy as np2import torch3from my3d import unproject4 5 6def subpixel_rays_from_img(H, W, K, c2w_pose, normalize_dir=True, f=8):7    assert c2w_pose[3, 3] == 1.8    H, W = H * f, W * f9    n = H * W10    ys, xs = np.meshgrid(range(H), range(W), indexing="ij")11    xy_coords = np.stack([xs, ys], axis=-1).reshape(n, 2)12 13    top_left = np.array([-0.5, -0.5]) + 1 / (2 * f)14    xy_coords = top_left + xy_coords / f15 16    ro = c2w_pose[:, -1]17    pts = unproject(K, xy_coords, depth=1)18    pts = pts @ c2w_pose.T19    rd = pts - ro20    rd = rd[:, :3]21    if normalize_dir:22        rd = rd / np.linalg.norm(rd, axis=-1, keepdims=True)23    ro = np.tile(ro[:3], (n, 1))24    return ro, rd25 26 27def rays_from_img(H, W, K, c2w_pose, normalize_dir=True):28    assert c2w_pose[3, 3] == 1.29    n = H * W30    ys, xs = np.meshgrid(range(H), range(W), indexing="ij")31    xy_coords = np.stack([xs, ys], axis=-1).reshape(n, 2)32 33    ro = c2w_pose[:, -1]34    pts = unproject(K, xy_coords, depth=1)35    pts = pts @ c2w_pose.T36    rd = pts - ro  # equivalently can subtract [0,0,0,1] before pose transform37    rd = rd[:, :3]38    if normalize_dir:39        rd = rd / np.linalg.norm(rd, axis=-1, keepdims=True)40    ro = np.tile(ro[:3], (n, 1))41    return ro, rd42 43 44def ray_box_intersect(ro, rd, aabb):45    """46    Intersection of ray with axis-aligned bounding box47    This routine works for arbitrary dimensions; commonly d = 2 or 348    only works for numpy, not torch (which has slightly diff api for min, max, and clone)49 50    Args:51        ro: [n, d] ray origin52        rd: [n, d] ray direction (assumed to be already normalized;53            if not still fine, meaning of t as time of flight holds true)54        aabb: [d, 2] bbox bound on each dim55    Return:56        is_intersect: [n,] of bool, whether the particular ray intersects the bbox57        t_min: [n,] ray entrance time58        t_max: [n,] ray exit time59    """60    n = ro.shape[0]61    d = aabb.shape[0]62    assert aabb.shape == (d, 2)63    assert ro.shape == (n, d) and rd.shape == (n, d)64 65    rd = rd.copy()66    rd[rd == 0] = 1e-6  # avoid div overflow; logically safe to give it big t67 68    ro = ro.reshape(n, d, 1)69    rd = rd.reshape(n, d, 1)70    ts = (aabb - ro) / rd  # [n, d, 2]71    t_min = ts.min(-1).max(-1)  # [n,] last of entrance72    t_max = ts.max(-1).min(-1)  # [n,] first of exit73    is_intersect = t_min < t_max74 75    return is_intersect, t_min, t_max76 77 78def as_torch_tsrs(device, *args):79    ret = []80    for elem in args:81        target_dtype = torch.float32 if np.issubdtype(elem.dtype, np.floating) else None82        ret.append(83            torch.as_tensor(elem, dtype=target_dtype, device=device)84        )85    return ret86 87 88def group_mask_filter(mask, *items):89    return [elem[mask] for elem in items]90 91 92def mask_back_fill(tsr, N, inds, base_value=1.0):93    shape = [N, *tsr.shape[1:]]94    canvas = base_value * np.ones_like(tsr, shape=shape)95    canvas[inds] = tsr96    return canvas97 98 99def render_one_view(model, aabb, H, W, K, pose):100    N = H * W101    bs = max(W * 5, 4096)  # render 5 rows; original batch size 4096, now 4000;102 103    ro, rd = rays_from_img(H, W, K, pose)104    ro, rd, t_min, t_max, intsct_inds = scene_box_filter(ro, rd, aabb)105    n = len(ro)106    # print(f"{n} vs {N}")  # n can be smaller than N since some rays do not intsct aabb107 108    # n = n // 1  # actual number of rays to render; only needed for fast debugging109 110    dev = model.device111    ro, rd, t_min, t_max = as_torch_tsrs(dev, ro, rd, t_min, t_max)112    rgbs = torch.zeros(n, 3, device=dev)113    depth = torch.zeros(n, 1, device=dev)114 115    with torch.no_grad():116        for i in range(int(np.ceil(n / bs))):117            s = i * bs118            e = min(n, s + bs)119            _rgbs, _depth, _ = render_ray_bundle(120                model, ro[s:e], rd[s:e], t_min[s:e], t_max[s:e]121            )122            rgbs[s:e] = _rgbs123            depth[s:e] = _depth124 125    rgbs, depth = rgbs.cpu().numpy(), depth.cpu().numpy()126 127    base_color = 1.0  # empty region needs to be white128    rgbs = mask_back_fill(rgbs, N, intsct_inds, base_color).reshape(H, W, 3)129    depth = mask_back_fill(depth, N, intsct_inds, base_color).reshape(H, W)130    return rgbs, depth131 132 133def scene_box_filter(ro, rd, aabb):134    N = len(ro)135 136    _, t_min, t_max = ray_box_intersect(ro, rd, aabb)137    # do not render what's behind the ray origin138    t_min, t_max = np.maximum(t_min, 0), np.maximum(t_max, 0)139    # can test intersect logic by reducing the focal length140    is_intsct = t_min < t_max141    ro, rd, t_min, t_max = group_mask_filter(is_intsct, ro, rd, t_min, t_max)142    intsct_inds = np.arange(N)[is_intsct]143    return ro, rd, t_min, t_max, intsct_inds144 145 146def render_ray_bundle(model, ro, rd, t_min, t_max):147    """148    The working shape is (k, n, 3) where k is num of samples per ray, n the ray batch size149    During integration the reduction is applied on k150 151    chain of filtering152    starting with ro, rd (from cameras), and a scene bbox153    - rays that do not intersect scene bbox; sample pts that fall outside the bbox154    - samples that do not fall within alpha mask155    - samples whose densities are very low; no need to compute colors on them156    """157    num_samples, step_size = model.get_num_samples((t_max - t_min).max())158    # print(num_samples)159    n, k = len(ro), num_samples160    # print(n,k)161    ticks = step_size * torch.arange(k, device=ro.device)162    ticks = ticks.view(k, 1, 1)163    t_min = t_min.view(n, 1)164    # t_min = t_min + step_size * torch.rand_like(t_min)  # NOTE seems useless165    t_max = t_max.view(n, 1)166    dists = t_min + ticks  # [n, 1], [k, 1, 1] -> [k, n, 1]167    pts = ro + rd * dists  # [n, 3], [n, 3], [k, n, 1] -> [k, n, 3]168    mask = (ticks < (t_max - t_min)).squeeze(-1)  # [k, 1, 1], [n, 1] -> [k, n, 1] -> [k, n]169    smp_pts = pts[mask]170 171    if model.alphaMask is not None:172        alphas = model.alphaMask.sample_alpha(smp_pts)173        alpha_mask = alphas > 0174        mask[mask.clone()] = alpha_mask175        smp_pts = pts[mask]176 177    σ = torch.zeros(k, n, device=ro.device)178    σ[mask] = model.compute_density_feats(smp_pts)179    weights = volume_rend_weights(σ, step_size)180    mask = weights > model.ray_march_weight_thres181    smp_pts = pts[mask]182 183    app_feats = model.compute_app_feats(smp_pts)184    # viewdirs = rd.view(1, n, 3).expand(k, n, 3)[mask]  # ray dirs for each point185    # additional wild factors here as in nerf-w; wild factors are optimizable186    c_dim = app_feats.shape[-1]187    colors = torch.zeros(k, n, c_dim, device=ro.device)188    colors[mask] = model.feats2color(app_feats)189 190    weights = weights.view(k, n, 1)  # can be used to compute other expected vals e.g. depth191    bg_weight = 1. - weights.sum(dim=0)  # [n, 1]192 193    rgbs = (weights * colors).sum(dim=0)  # [n, 3]194 195    if model.blend_bg_texture:196        uv = spherical_xyz_to_uv(rd)197        bg_feats = model.compute_bg(uv)198        bg_color = model.feats2color(bg_feats)199        rgbs = rgbs + bg_weight * bg_color200    else:201        rgbs = rgbs + bg_weight * 1.  # blend white bg color202    # print(rgbs.shape)203    # rgbs = rgbs.clamp(0, 1)  # don't clamp since this is can be SD latent features204 205    E_dists = (weights * dists).sum(dim=0)206    bg_dist = 10.  # blend bg distance; just don't make it too large207    E_dists = E_dists + bg_weight * bg_dist208    return rgbs, E_dists, weights.squeeze(-1)209 210 211def spherical_xyz_to_uv(xyz):212    # xyz is Tensor of shape [N, 3], uv in [-1, 1]213    x, y, z = xyz.t()  # [N]214    xy = (x ** 2 + y ** 2) ** 0.5215    u = torch.atan2(xy, z) / torch.pi  # [N]216    v = torch.atan2(y, x) / (torch.pi * 2) + 0.5  # [N]217    uv = torch.stack([u, v], -1)  # [N, 2]218    uv = uv * 2 - 1  # [0, 1] -> [-1, 1]219    return uv220 221 222def volume_rend_weights(σ, dist):223    α = 1 - torch.exp(-σ * dist)224    T = torch.ones_like(α)225    T[1:] = (1 - α).cumprod(dim=0)[:-1]226    assert (T >= 0).all()227    weights = α * T228    return weights229