"""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 # ---------------------------------------------------------------- preprocessing 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() # uint8 DIS alpha values below this floor are set to 0 before the author's preprocessing (0 = off). DIS leaves faint # responses BiRefNet does not (1-5 of 255 at the image corners, along thin wisps); the author's rule crops to the # nonzero alpha, so they stretch the crop. The value is calibrated against BiRefNet's crop on held-out images # (scripts/calibrate_matting_floor.py, contract/matting_floor.json). DIS_ALPHA_FLOOR = 8 # calibrated 2026-11-09 on 31 held-out gallery images: crop box error mean 0.14 %, max 0.59 % of the side 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) # ---------------------------------------------------------------- shared randomness 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) # ---------------------------------------------------------------- octree sampling 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] # ---------------------------------------------------------------- Gaussians and PLY 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='