Spaces:
Running on Zero
Running on Zero
| # Modified for SplatAnything: inference-only release. | |
| """Image preprocessing and PLY/depth export.""" | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| import PIL.Image | |
| import torch | |
| from src.utils import cropping | |
| from src.utils.camera_io import list_image_files, load_cameras_json | |
| from src.utils.sky_segmentation import SkySegmenter | |
| from src.utils.save_utils import save_gs_ply, save_scene_ply, save_image_png, save_depth_png, save_depth_npy | |
| def prepare_input(input_path): | |
| folder = Path(input_path) | |
| images = [str(p) for p in list_image_files(folder)] | |
| return (images, folder.name) | |
| def _centered_default_intrinsics(image_size): | |
| """Dummy centered intrinsics used only to reproduce camera-free image preprocessing. | |
| In RGB-only mode this matrix is never passed to the model or renderer. It only lets the | |
| camera-aware crop helper reduce to an aspect-preserving center crop. | |
| """ | |
| (width, height) = image_size | |
| focal = float(max(width, height)) | |
| return np.array([[focal, 0.0, width / 2.0], [0.0, focal, height / 2.0], [0.0, 0.0, 1.0]], dtype=np.float32) | |
| def preprocess_view(image: 'PIL.Image.Image', K: np.ndarray=None, resolution_wh=(518, 518), depthmap=None): | |
| image = image.convert('RGB') | |
| K = _centered_default_intrinsics(image.size) if K is None else K.astype(np.float32).copy() | |
| if not (0 < K[0, 2] < image.width and 0 < K[1, 2] < image.height): | |
| raise ValueError('Camera principal point must lie inside the original image.') | |
| if depthmap is not None: | |
| depthmap = np.ascontiguousarray(depthmap, dtype=np.float32) | |
| if depthmap.shape[:2] != image.size[::-1]: | |
| depthmap = cv2.resize(depthmap, image.size, interpolation=cv2.INTER_NEAREST) | |
| (W, H) = image.size | |
| (cx, cy) = K[:2, 2].round().astype(int) | |
| (mx, my) = (min(cx, W - cx), min(cy, H - cy)) | |
| (image, depthmap, K) = cropping.crop_image_depthmap(image, depthmap, K, (cx - mx, cy - my, cx + mx, cy + my)) | |
| res = np.array(resolution_wh) | |
| (image, depthmap, K) = cropping.rescale_image_depthmap(image, depthmap, K, res) | |
| K2 = cropping.camera_matrix_of_crop(K, image.size, resolution_wh, offset_factor=0.5) | |
| crop_bbox = cropping.bbox_from_intrinsics_in_out(K, K2, resolution_wh) | |
| (image, depthmap, K2) = cropping.crop_image_depthmap(image, depthmap, K, crop_bbox) | |
| img_t = torch.from_numpy(np.array(image, dtype=np.float32).transpose(2, 0, 1).copy()) / 255.0 | |
| K_t = torch.from_numpy(K2.astype(np.float32)) | |
| if depthmap is None: | |
| return (img_t, K_t) | |
| return (img_t, K_t, torch.from_numpy(np.ascontiguousarray(depthmap, dtype=np.float32))) | |
| def build_views(input_path: str, cameras_path: str=None, resolution_wh=(518, 518), device='cpu', cam_align: bool=True, max_images: int=None, sky_segmenter: 'SkySegmenter'=None, sky_max_fraction: float=0.95): | |
| (img_paths, scene_name) = prepare_input(input_path) | |
| if scene_name in ('images', 'image', 'imgs') and Path(input_path).parent.name: | |
| scene_name = Path(input_path).parent.name | |
| if max_images is not None and max_images > 0: | |
| img_paths = img_paths[:max_images] | |
| print(f'[Input] Capped to first {len(img_paths)} images (--max_images={max_images})') | |
| has_camera_prior = cameras_path is not None | |
| if has_camera_prior: | |
| if not Path(cameras_path).is_file(): | |
| raise FileNotFoundError(f'Camera JSON not found: {cameras_path}') | |
| (extr_map, intr_map) = load_cameras_json(cameras_path) | |
| else: | |
| (extr_map, intr_map) = ({}, {}) | |
| (imgs, Ks, c2ws, depths, sky_masks) = ([], [], [], [], []) | |
| for p in img_paths: | |
| stem = Path(p).stem | |
| if has_camera_prior and (stem not in extr_map or stem not in intr_map): | |
| raise ValueError(f"No camera entry (extrinsics/intrinsics) for image '{stem}' in {cameras_path}.") | |
| prior_d = None | |
| K = intr_map[stem] if has_camera_prior else None | |
| out = preprocess_view(PIL.Image.open(p), K, resolution_wh, depthmap=prior_d) | |
| if prior_d is None: | |
| (img_t, K2) = out | |
| else: | |
| (img_t, K2, depth_t) = out | |
| depths.append(depth_t) | |
| imgs.append(img_t) | |
| if has_camera_prior: | |
| Ks.append(K2) | |
| c2ws.append(torch.from_numpy(extr_map[stem])) | |
| if sky_segmenter is not None: | |
| rgb_hw3 = (img_t.permute(1, 2, 0).numpy() * 255.0).astype(np.uint8) | |
| sky = sky_segmenter.sky_mask(rgb_hw3) | |
| frac = float(sky.mean()) | |
| if frac > sky_max_fraction: | |
| print(f"[SkySeg] '{stem}': sky fraction {frac:.2%} > {sky_max_fraction:.0%} -> leaving view UNFILTERED (likely false positive).") | |
| sky = np.zeros_like(sky) | |
| sky_masks.append(torch.from_numpy(sky)) | |
| imgs = torch.stack(imgs).unsqueeze(0).to(device) | |
| views = {'img': imgs} | |
| if has_camera_prior: | |
| intr = torch.stack(Ks).unsqueeze(0).to(device).float() | |
| extr = torch.stack(c2ws).unsqueeze(0).to(device).float() | |
| if cam_align: | |
| first_inv = torch.linalg.inv(extr[:, 0]) | |
| extr = first_inv.unsqueeze(1) @ extr | |
| views['camera_poses'] = extr | |
| views['camera_intrs'] = intr | |
| if depths: | |
| views['depthmap'] = torch.stack(depths).unsqueeze(0).to(device).float() | |
| print(f"[Input] Loaded prior_depth for {len(depths)} views (sky pixels depth<0: {(views['depthmap'] < 0).sum().item()})") | |
| if sky_masks: | |
| sky = torch.stack(sky_masks).unsqueeze(0).to(device) | |
| if 'depthmap' not in views: | |
| views['depthmap'] = torch.zeros_like(sky, dtype=torch.float32) | |
| views['depthmap'] = views['depthmap'].masked_fill(sky, -1.0) | |
| print(f'[SkySeg] tagged {int(sky.sum().item())} sky pixels across {len(sky_masks)} views ({100.0 * sky.float().mean().item():.1f}% overall).') | |
| return (views, scene_name, img_paths) | |
| def _batch0(splats, key): | |
| """Pull batch 0 of a splat field (list-per-batch after prune_gs, or [B, ...] before).""" | |
| v = splats.get(key) | |
| if v is None: | |
| return None | |
| return v[0] | |
| def export_gaussians(preds, outdir: Path, opacity_thresh: float, scale_quantile): | |
| """Write the pruned Gaussian splats (batch 0) to gaussians.ply. | |
| SH DC is stored raw into f_dc (the viewer reconstructs rgb = 0.5 + C0 * f_dc), matching | |
| save_validation_gs_ply. Near-transparent Gaussians below opacity_thresh are dropped. | |
| """ | |
| splats = preds.get('splats') | |
| if splats is None or 'means' not in splats: | |
| print(' [gaussians] no splats in predictions -> skip') | |
| return | |
| means = _batch0(splats, 'means') | |
| scales = _batch0(splats, 'scales') | |
| quats = _batch0(splats, 'quats') | |
| opacities = _batch0(splats, 'opacities') | |
| sh = _batch0(splats, 'sh') | |
| if means is None or means.shape[0] == 0: | |
| print(' [gaussians] empty cloud -> skip') | |
| return | |
| dc = sh[:, 0, :] if sh is not None else torch.zeros_like(means) | |
| if opacity_thresh and opacities is not None: | |
| keep = opacities.reshape(-1) >= opacity_thresh | |
| if keep.any(): | |
| (means, scales, quats) = (means[keep], scales[keep], quats[keep]) | |
| (dc, opacities) = (dc[keep], opacities.reshape(-1)[keep]) | |
| else: | |
| print(' [gaussians] all Gaussians below opacity threshold -> skip') | |
| return | |
| path = outdir / 'gaussians.ply' | |
| save_gs_ply(path, means, scales, quats, dc, opacities, quantile_threshold=scale_quantile) | |
| print(f' [gaussians] saved {means.shape[0]} Gaussians -> {path}') | |
| def export_depth_rgb(preds, views, img_paths, outdir: Path): | |
| depth = preds.get('depth') | |
| imgs = views.get('img') | |
| if depth is None: | |
| print(' [depth/rgb] no predicted depth -> skip') | |
| return | |
| depth = depth[0] | |
| imgs = imgs[0] | |
| S = depth.shape[0] | |
| depth_dir = outdir / 'depth' | |
| rgb_dir = outdir / 'rgb' | |
| depth_dir.mkdir(parents=True, exist_ok=True) | |
| rgb_dir.mkdir(parents=True, exist_ok=True) | |
| for i in range(S): | |
| stem = Path(img_paths[i]).stem if i < len(img_paths) else f'{i:04d}' | |
| d = depth[i, ..., 0].float() | |
| save_depth_png(depth_dir / f'{stem}.png', d) | |
| save_depth_npy(depth_dir / f'{stem}.npy', d) | |
| save_image_png(rgb_dir / f'{stem}.png', imgs[i].permute(1, 2, 0).float()) | |
| print(f' [depth/rgb] saved {S} predicted depth maps -> {depth_dir} and input RGB -> {rgb_dir}') | |
| def export_points(preds, outdir): | |
| (means, rgb) = (preds['points_ply_means'][0], preds['points_ply_rgb'][0]) | |
| colors = (rgb.clamp(0, 1) * 255).round().to(torch.uint8) | |
| save_scene_ply(outdir / 'points.ply', means, colors) | |
| print(f' [points] saved {len(means)} points') | |