| """Host-side steps of the Idol pipeline that are not tensor programs, as the phone runs them (the Kotlin code in the |
| sample mirrors these functions line by line). Each follows idol-sdk @4702a7f; the comments name the source. |
| |
| - preprocess: idol/preprocess.py preprocess_image, with the alpha predictor passed in (BiRefNet or the matting graph) |
| - shared randomness: CONTRACT.md "Shared randomness" |
| - systematic_rows / octree_descend: idol/models/decoder.py sample + systematic_sample, with the per-row pool draws |
| - to_gaussians / write_ply: GaussianFixedlenDecoder.to_representation, idol/gaussian.py Gaussian + save_ply |
| """ |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from PIL import Image, ImageFilter |
|
|
| CROP_MARGIN = {'synthetic': 1.2, 'real': 1.7} |
| IMAGE_SIZE = 1024 |
| MAX_LEVEL = 8 |
| PER_ANCHOR = 32 |
| PHASE_THRESHOLD = 0.5 |
|
|
|
|
| |
|
|
| def has_real_alpha(image): |
| return image.mode == 'RGBA' and np.array(image.getchannel(3), dtype=np.int32).min() < 255 |
|
|
|
|
| def preprocess(image, alpha_fn, image_size=IMAGE_SIZE, crop_margin=1.2): |
| """-> (composite RGB image the encoders see, uint8 alpha at the resized size before erosion, crop box). |
| alpha_fn(rgb PIL image) -> uint8 [H, W] alpha.""" |
| width, height = image.size |
| scale = image_size / min(width, height) |
| image = image.resize((max(1, int(round(width * scale))), max(1, int(round(height * scale)))), Image.LANCZOS) |
| if not has_real_alpha(image): |
| rgb = image.convert('RGB') |
| image = rgb.copy() |
| image.putalpha(Image.fromarray(alpha_fn(rgb), mode='L')) |
| alpha_raw = np.array(image.getchannel(3)) |
| image.putalpha(image.getchannel(3).filter(ImageFilter.MinFilter(3))) |
| alpha = np.array(image.getchannel(3)) |
| ys, xs = np.nonzero(alpha) |
| if len(xs) == 0: |
| raise ValueError('no object found in the image (the alpha channel is empty)') |
| left, top, right, bottom = xs.min(), ys.min(), xs.max(), ys.max() |
| center_x, center_y = (left + right) / 2, (top + bottom) / 2 |
| half = max(right - left, bottom - top) / 2 * crop_margin |
| box = [int(center_x - half), int(center_y - half), int(center_x + half), int(center_y + half)] |
| image = image.crop(box) |
| image = image.resize((image_size, image_size), Image.LANCZOS) |
| composite = Image.new('RGB', (image_size, image_size), (0, 0, 0)) |
| composite.paste(image, mask=image.split()[3]) |
| return composite, alpha_raw, box |
|
|
|
|
| def birefnet_alpha_fn(birefnet): |
| return lambda rgb: np.array(birefnet.remove_background(rgb).getchannel(3)) |
|
|
|
|
| BIREFNET_SIZE = 1024 |
| BIREFNET_MEAN = (0.485, 0.456, 0.406) |
| BIREFNET_STD = (0.229, 0.224, 0.225) |
|
|
|
|
| def birefnet_alpha(rgb, run_bir): |
| """The author's BiRefNet.remove_background (idol/models/birefnet.py) around the graph: RGB / 255, bilinear |
| (align_corners) to 1024 x 1024, ImageNet mean / std; the graph returns the sigmoid map [1, 1, 1024, 1024]; |
| bilinear (align_corners) back to the image size; x 255 truncated to uint8.""" |
| w, h = rgb.size |
| t = torch.from_numpy(np.array(rgb.convert('RGB'), dtype=np.float32) / 255.0).permute(2, 0, 1)[None] |
| t = F.interpolate(t, size=(BIREFNET_SIZE, BIREFNET_SIZE), mode='bilinear', align_corners=True) |
| t = (t - torch.tensor(BIREFNET_MEAN).view(1, 3, 1, 1)) / torch.tensor(BIREFNET_STD).view(1, 3, 1, 1) |
| a = torch.from_numpy(np.asarray(run_bir(t.numpy()), np.float32).reshape(1, 1, BIREFNET_SIZE, BIREFNET_SIZE)) |
| a = F.interpolate(a, size=(h, w), mode='bilinear', align_corners=True)[0, 0] |
| return (a.clamp(0, 1) * 255).to(torch.uint8).numpy() |
|
|
|
|
| |
| |
| |
| |
| DIS_ALPHA_FLOOR = 8 |
|
|
|
|
| def dis_alpha(rgb, run_dis, floor=None): |
| """IS-Net DIS general-use as its authors run it (xuebinqin/DIS inference): bilinear resize to 1024², x/255 - 0.5, |
| the sigmoid map back to the image size (bilinear), min-max normalised, x255 truncated to uint8. |
| run_dis(x [1, 3, 1024, 1024] float32 numpy) -> [1, 1, 1024, 1024] sigmoid alpha.""" |
| w, h = rgb.size |
| t = torch.from_numpy(np.array(rgb, dtype=np.float32)).permute(2, 0, 1)[None] |
| t = F.interpolate(t, size=(1024, 1024), mode='bilinear', align_corners=False) |
| x = (t / 255.0 - 0.5).numpy().astype(np.float32) |
| a = torch.from_numpy(np.asarray(run_dis(x), dtype=np.float32).reshape(1, 1, 1024, 1024)) |
| a = F.interpolate(a, size=(h, w), mode='bilinear', align_corners=False)[0, 0] |
| mi, ma = a.min(), a.max() |
| a = (a - mi) / (ma - mi).clamp_min(1e-12) |
| a = (a.clamp(0, 1) * 255).to(torch.uint8).numpy() |
| floor = DIS_ALPHA_FLOOR if floor is None else floor |
| if floor > 0: |
| a = np.where(a >= floor, a, 0).astype(np.uint8) |
| return a |
|
|
|
|
| def pixels_of(composite): |
| """The encoders' input: [1, 3, 1024, 1024] float32 in [0, 1].""" |
| return (np.asarray(composite, dtype=np.float32).transpose(2, 0, 1)[None] / 255.0).astype(np.float32) |
|
|
|
|
| |
|
|
| def make_noise(seed, q_tokens=8192, in_ch=16, cam_ch=5): |
| g = torch.Generator().manual_seed(seed) |
| vae = torch.randn(1, 32, 128, 128, generator=g) |
| latent = torch.randn(1, q_tokens, in_ch, generator=g) |
| camera = torch.randn(1, 1, cam_ch, generator=g) |
| g1 = torch.Generator().manual_seed(seed + 1) |
| pool = torch.rand(MAX_LEVEL, q_tokens, generator=g1) |
| leaf = torch.rand(1, q_tokens, 3, generator=g1) |
| return {k: v.numpy() for k, v in |
| dict(vae_noise=vae, latent_noise=latent, camera_noise=camera, octree_pool=pool, leaf_pool=leaf).items()} |
|
|
|
|
| def time_grid(steps=5, rescale_t=3.0): |
| t = np.linspace(1, 0, steps + 1) |
| return rescale_t * t / (1 + (rescale_t - 1) * t) |
|
|
|
|
| |
|
|
| def systematic_rows(probs, counts, u): |
| """idol systematic_sample for one batch, row i's offset u0 = u[i] / n_i. probs [Np, P] float32 torch, counts [Np] |
| int64 torch, u [Np] float32 torch -> [Np, P] int64.""" |
| P = probs.size(-1) |
| probs = probs.to(torch.float32).clamp_min(0) |
| row_sums = probs.sum(1, keepdim=True) |
| zero_mask = row_sums.eq(0) |
| probs = probs / row_sums.clamp_min(1) |
| if zero_mask.any(): |
| probs = probs.clone() |
| probs[zero_mask.expand_as(probs)] = 1.0 / P |
| out = torch.zeros(probs.shape[0], P, dtype=torch.long) |
| cdf = probs.cumsum(dim=1).clamp(max=1.0 - 1e-12) |
| for n in torch.unique(counts).tolist(): |
| if n == 0: |
| continue |
| rows = (counts == n).nonzero(as_tuple=False).squeeze(1) |
| u0 = u.index_select(0, rows)[:, None] / float(n) |
| grid = torch.arange(n, dtype=torch.float32)[None, :] / float(n) |
| us = (u0 + grid).clamp(max=1.0 - 1e-12) |
| idx = torch.searchsorted(cdf.index_select(0, rows), us).clamp_max(P - 1) |
| buf = torch.zeros(rows.numel(), P, dtype=torch.float32) |
| buf.scatter_add_(1, idx, torch.ones_like(idx, dtype=buf.dtype)) |
| out.index_copy_(0, rows, buf.to(torch.long)) |
| return out |
|
|
|
|
| CHILD_OFFSET = torch.tensor([[i, j, k] for k in [0, 1] for j in [0, 1] for i in [0, 1]], dtype=torch.long) |
|
|
|
|
| def octree_descend(probs_fn, pool, leaf, num_points=8192, level=MAX_LEVEL, rec=None): |
| """idol OctreeProbabilityFixedlenDecoder.sample for one batch. probs_fn(cells [1, L, 3] float32 torch, l) -> |
| child probabilities [1, L, 8] (the softmax of the decoder's logits). pool [8, >= rows] and leaf [1, num_points, 3] |
| float32 torch. -> anchors [1, num_points, 3] in [0, 1]^3.""" |
| prev = torch.zeros(1, 3, dtype=torch.long) |
| counts = torch.full((1,), num_points, dtype=torch.long) |
| for l in range(1, level + 1): |
| res_p = 1 << (l - 1) |
| cells = (prev.to(torch.float32) + 0.5) / res_p |
| probs = probs_fn(cells[None], l)[0] |
| sampled = systematic_rows(probs, counts, pool[l - 1, :prev.shape[0]]).flatten() |
| child = (prev[:, None, :] * 2 + CHILD_OFFSET[None]).flatten(0, 1) |
| mask = sampled > 0 |
| if rec is not None: |
| rec[f'oct_cells_{l}'] = cells.numpy() |
| rec[f'oct_probs_{l}'] = probs.float().numpy() |
| rec[f'oct_counts_{l}'] = counts.numpy() |
| prev, counts = child[mask], sampled[mask] |
| coords = torch.repeat_interleave(prev, counts, dim=0) |
| return ((coords.to(torch.float32) + leaf[0]) / (1 << level))[None] |
|
|
|
|
| |
|
|
| def inverse_sigmoid(x): |
| return torch.log(x / (1 - x)) |
|
|
|
|
| def inverse_softplus(x): |
| return x + torch.log(-torch.expm1(-x)) |
|
|
|
|
| LAYOUT = [('_xyz', 3), ('_features_dc', 3), ('_scaling', 3), ('_rotation', 4), ('_opacity', 1), ('_offset_scale', 1)] |
|
|
|
|
| def to_gaussians(points, features, rep, perturbation, base_offset_scale): |
| """GaussianFixedlenDecoder.to_representation + Gaussian.__init__ for one batch, float32. |
| points [1, L, 3], features [1, L, 480] torch -> dict of activated tensors (z-up): xyz [N, 3], features_dc [N, 1, 3], |
| scaling [N, 3], rotation [N, 4] (w, x, y, z), opacity [N, 1].""" |
| features = features.to(torch.float32) |
| L = features.shape[1] |
| sl, start = {}, 0 |
| for name, size in LAYOUT: |
| sl[name] = features[0, :, start:start + PER_ANCHOR * size].reshape(L, PER_ANCHOR, size) |
| start += PER_ANCHOR * size |
| lr = rep['lr'] |
| offset_scale = F.softplus(sl['_offset_scale'] + base_offset_scale) |
| offset = sl['_xyz'] * lr['_xyz'] + perturbation |
| offset = torch.tanh(offset) * 0.5 * rep['perturbe_size'] * offset_scale |
| raw_xyz = (offset + points[0, :, None, :]).flatten(0, 1) |
| scaling_offset = inverse_softplus(torch.tensor(rep['scaling_bias'])) |
| opacity_offset = inverse_sigmoid(torch.tensor(rep['opacity_bias'])) |
| scaling = F.softplus(sl['_scaling'].flatten(0, 1) * lr['_scaling'] + scaling_offset) |
| return { |
| 'xyz': raw_xyz * 1.0 - 0.5, |
| 'features_dc': (sl['_features_dc'].flatten(0, 1) * lr['_features_dc'])[:, None, :], |
| 'scaling': torch.sqrt(torch.square(scaling) + rep['3d_filter_kernel_size'] ** 2), |
| 'rotation': F.normalize(sl['_rotation'].flatten(0, 1) * lr['_rotation'] + torch.tensor([1.0, 0.0, 0.0, 0.0])), |
| 'opacity': torch.sigmoid(sl['_opacity'].flatten(0, 1) * lr['_opacity'] + opacity_offset), |
| } |
|
|
|
|
| def quaternion_multiply(a, b): |
| aw, ax, ay, az = a |
| bw, bx, by, bz = b[:, 0], b[:, 1], b[:, 2], b[:, 3] |
| return np.stack([aw * bw - ax * bx - ay * by - az * bz, aw * bx + ax * bw + ay * bz - az * by, |
| aw * by - ax * bz + ay * bw + az * bx, aw * bz + ax * by - ay * bx + az * bw], axis=1) |
|
|
|
|
| PLY_NAMES = ['x', 'y', 'z', 'nx', 'ny', 'nz', 'f_dc_0', 'f_dc_1', 'f_dc_2', 'opacity', |
| 'scale_0', 'scale_1', 'scale_2', 'rot_0', 'rot_1', 'rot_2', 'rot_3'] |
|
|
|
|
| def ply_rows(g): |
| """idol Gaussian.save_ply's attribute rows [N, 17] float32 (z-up turned to y-up).""" |
| xyz = g['xyz'].detach().cpu().numpy() |
| f_dc = g['features_dc'].detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() |
| eps = 1e-6 |
| opacities = inverse_sigmoid(g['opacity'].clamp(eps, 1 - eps)).detach().cpu().numpy() |
| scale = torch.log(g['scaling'].clamp_min(eps)).detach().cpu().numpy() |
| rotation = g['rotation'].detach().cpu().numpy() |
| xyz = np.stack([xyz[:, 0], -xyz[:, 2], xyz[:, 1]], axis=1) |
| rotation = quaternion_multiply(np.array([np.sqrt(0.5), np.sqrt(0.5), 0.0, 0.0]), rotation) |
| return np.concatenate((xyz, np.zeros_like(xyz), f_dc, opacities, scale, rotation), axis=1).astype(np.float32) |
|
|
|
|
| def write_ply(rows, path): |
| """Binary little-endian PLY with the 17 float properties of save_ply (the same bytes plyfile writes).""" |
| header = ['ply', 'format binary_little_endian 1.0', f'element vertex {rows.shape[0]}'] |
| header += [f'property float {n}' for n in PLY_NAMES] + ['end_header'] |
| with open(path, 'wb') as f: |
| f.write(('\n'.join(header) + '\n').encode('ascii')) |
| f.write(np.ascontiguousarray(rows, dtype='<f4').tobytes()) |
|
|
|
|
| def read_ply(path): |
| with open(path, 'rb') as f: |
| data = f.read() |
| end = data.index(b'end_header\n') + len(b'end_header\n') |
| head = data[:end].decode('ascii').splitlines() |
| names = [l.split()[-1] for l in head if l.startswith('property')] |
| n = int([l for l in head if l.startswith('element vertex')][0].split()[-1]) |
| return names, np.frombuffer(data[end:], dtype='<f4').reshape(n, len(names)) |
|
|
