import os import cv2 import torch import shutil from datetime import datetime import glob import gc import spaces import gradio as gr import numpy as np import concerto from scipy.spatial import cKDTree from scipy.spatial.transform import Rotation as R import trimesh import time from typing import List, Tuple from pathlib import Path from einops import rearrange from tqdm import tqdm from PIL import Image from torchvision import transforms as TF try: import flash_attn except ImportError: flash_attn = None from visual_util import predictions_to_glb from vggt.models.vggt import VGGT from vggt.utils.load_fn import load_and_preprocess_images from vggt.utils.pose_enc import pose_encoding_to_extri_intri from vggt.utils.geometry import unproject_depth_map_to_point_map # set random seed # (random seed affect pca color, yet change random seed need manual adjustment kmeans) # (the pca prevent in paper is with another version of cuda and pytorch environment) concerto.utils.set_seed(53124) # Load model (to CPU; moved to GPU on-demand via @spaces.GPU) if flash_attn is not None: print("Loading model with Flash Attention.") concerto_model = concerto.load("concerto_large", repo_id="Pointcept/Concerto") sonata_model = concerto.model.load("sonata", repo_id="facebook/sonata") else: print("Loading model without Flash Attention.") custom_config = dict( # enc_patch_size=[1024 for _ in range(5)], # reduce patch size if necessary enable_flash=False, ) concerto_model = concerto.load( "concerto_large", repo_id="Pointcept/Concerto", custom_config=custom_config ) sonata_model = concerto.load("sonata", repo_id="facebook/sonata", custom_config=custom_config) transform = concerto.transform.default() VGGT_model = VGGT() _URL = "https://huggingface.co/facebook/VGGT-1B/resolve/main/model.pt" VGGT_model.load_state_dict(torch.hub.load_state_dict_from_url(_URL)) def _estimate_normals_knn(points, k=30): """Estimate per-point normals via PCA over k nearest neighbors.""" points = np.asarray(points, dtype=np.float64) n = len(points) if n == 0: return np.zeros((0, 3), dtype=np.float64) k = min(k, n) tree = cKDTree(points) _, idx = tree.query(points, k=k) neighbors = points[idx] # (n, k, 3) centered = neighbors - neighbors.mean(axis=1, keepdims=True) cov = np.einsum("nki,nkj->nij", centered, centered) / max(k - 1, 1) # eigh returns ascending eigenvalues; smallest -> normal direction. _, eigvecs = np.linalg.eigh(cov) normals = eigvecs[:, :, 0] norms = np.linalg.norm(normals, axis=1, keepdims=True) norms = np.where(norms < 1e-12, 1.0, norms) return normals / norms def _segment_plane_ransac(points, distance_threshold, ransac_n=3, num_iterations=1000): """RANSAC plane fitting. Returns (plane_model=[a,b,c,d], inlier_indices).""" points = np.asarray(points, dtype=np.float64) n = len(points) if n < ransac_n: raise ValueError("Not enough points for plane fitting.") rng = np.random.default_rng() best_inliers = np.empty(0, dtype=np.int64) best_plane = None for _ in range(num_iterations): sample_idx = rng.choice(n, size=ransac_n, replace=False) sample = points[sample_idx] v1 = sample[1] - sample[0] v2 = sample[2] - sample[0] normal = np.cross(v1, v2) nrm = np.linalg.norm(normal) if nrm < 1e-12: continue normal = normal / nrm d = -np.dot(normal, sample[0]) dist = np.abs(points @ normal + d) inliers = np.where(dist <= distance_threshold)[0] if len(inliers) > len(best_inliers): best_inliers = inliers best_plane = np.array([normal[0], normal[1], normal[2], d], dtype=np.float64) if best_plane is None: raise ValueError("RANSAC failed to find any valid plane.") # Refit plane to inliers via SVD for a tighter estimate. inlier_pts = points[best_inliers] centroid = inlier_pts.mean(axis=0) _, _, vh = np.linalg.svd(inlier_pts - centroid, full_matrices=False) normal = vh[-1] nrm = np.linalg.norm(normal) if nrm > 1e-12: normal = normal / nrm d = -np.dot(normal, centroid) dist = np.abs(points @ normal + d) refined = np.where(dist <= distance_threshold)[0] if len(refined) >= len(best_inliers): best_inliers = refined best_plane = np.array([normal[0], normal[1], normal[2], d], dtype=np.float64) return best_plane, best_inliers def _read_ply(path): """Read a PLY file via trimesh. Returns (points, colors[0..1], normals|None).""" obj = trimesh.load(path, process=False) if isinstance(obj, trimesh.PointCloud): points = np.asarray(obj.vertices, dtype=np.float64) colors = None if obj.colors is not None and len(obj.colors) == len(points): colors = np.asarray(obj.colors, dtype=np.float64)[:, :3] / 255.0 normals = None # trimesh.PointCloud may carry vertex_normals via metadata. meta = getattr(obj, "metadata", {}) or {} vn = meta.get("ply_raw", {}).get("vertex", {}).get("data", None) if meta else None if vn is not None and "nx" in vn.dtype.names: normals = np.stack([vn["nx"], vn["ny"], vn["nz"]], axis=-1).astype(np.float64) elif isinstance(obj, trimesh.Trimesh): points = np.asarray(obj.vertices, dtype=np.float64) if obj.visual is not None and hasattr(obj.visual, "vertex_colors") and obj.visual.vertex_colors is not None: colors = np.asarray(obj.visual.vertex_colors, dtype=np.float64)[:, :3] / 255.0 else: colors = None normals = np.asarray(obj.vertex_normals, dtype=np.float64) if obj.vertex_normals is not None else None else: raise ValueError(f"Unsupported PLY content type: {type(obj)}") if colors is None: colors = np.ones_like(points) return points, colors, normals def pad_0001(Ts): """Pad (N,3,4) or (3,4) extrinsics with [0,0,0,1] row to become (N,4,4)/(4,4).""" Ts = np.asarray(Ts, dtype=np.float64) if Ts.ndim == 2: if Ts.shape == (4, 4): return Ts if Ts.shape == (3, 4): return np.vstack([Ts, np.array([[0.0, 0.0, 0.0, 1.0]])]) raise ValueError(f"Unexpected T shape: {Ts.shape}") if Ts.ndim == 3: if Ts.shape[1:] == (4, 4): return Ts if Ts.shape[1:] == (3, 4): n = Ts.shape[0] bottom = np.tile(np.array([[[0.0, 0.0, 0.0, 1.0]]]), (n, 1, 1)) return np.concatenate([Ts, bottom], axis=1) raise ValueError(f"Unexpected T shape: {Ts.shape}") raise ValueError(f"Unexpected T ndim: {Ts.ndim}") def T_to_C(T): """Convert world->cam extrinsic to camera center C = -R^T @ t.""" T = np.asarray(T, dtype=np.float64) R_mat = T[:3, :3] t = T[:3, 3] return -R_mat.T @ t def im_distance_to_im_depth(im_dist, K): """Convert per-pixel ray distance to per-pixel depth (Z in camera frame).""" im_dist = np.asarray(im_dist, dtype=np.float64) H, W = im_dist.shape fx, fy = K[0, 0], K[1, 1] cx, cy = K[0, 2], K[1, 2] us = np.arange(W) vs = np.arange(H) uu, vv = np.meshgrid(us, vs) # (H, W) x = (uu - cx) / fx y = (vv - cy) / fy norm = np.sqrt(x * x + y * y + 1.0) return im_dist / norm def im_depth_to_point_cloud(im_depth, K, T): """ Backproject a depth image to world points. K: (3,3) intrinsics. T: (4,4) world->cam extrinsic. Returns (H*W, 3) world coordinates (no invalid filtering, matches the `to_image=False, ignore_invalid=False` semantics that the call sites use). """ im_depth = np.asarray(im_depth, dtype=np.float64) H, W = im_depth.shape fx, fy = K[0, 0], K[1, 1] cx, cy = K[0, 2], K[1, 2] us = np.arange(W) vs = np.arange(H) uu, vv = np.meshgrid(us, vs) z = im_depth x = (uu - cx) * z / fx y = (vv - cy) * z / fy pts_cam = np.stack([x, y, z], axis=-1).reshape(-1, 3) # (H*W, 3) R_mat = T[:3, :3] t = T[:3, 3] # world point: P_w = R^T (P_c - t) pts_world = (pts_cam - t) @ R_mat return pts_world def axis_angle_to_matrix(axis_angle): """axis_angle: (3,) vector whose direction is axis and magnitude is angle (rad).""" return R.from_rotvec(np.asarray(axis_angle, dtype=np.float64)).as_matrix() class PointCloud: """Minimal replacement for o3d.geometry.PointCloud used by this app.""" def __init__(self, points=None, colors=None, normals=None): self.points = None if points is None else np.asarray(points, dtype=np.float64) self.colors = None if colors is None else np.asarray(colors, dtype=np.float64) self.normals = None if normals is None else np.asarray(normals, dtype=np.float64) def has_colors(self): return self.colors is not None and len(self.colors) > 0 def has_normals(self): return self.normals is not None and len(self.normals) > 0 def select_by_index(self, idx): idx = np.asarray(idx, dtype=np.int64) return PointCloud( points=self.points[idx], colors=self.colors[idx] if self.has_colors() else None, normals=self.normals[idx] if self.has_normals() else None, ) def estimate_normals(self, k=30): self.normals = _estimate_normals_knn(self.points, k=k) @spaces.GPU def _gpu_run_vggt_inference(images_tensor): """ GPU-only function: Run VGGT model inference on preprocessed images. """ global VGGT_model device = "cuda" if torch.cuda.is_available() else "cpu" images_tensor = images_tensor.to(device) model = VGGT_model.to(device) model.eval() print("Running inference...") with torch.no_grad(): if device == "cuda": with torch.cuda.amp.autocast(dtype=torch.bfloat16): predictions = model(images_tensor) else: predictions = model(images_tensor) print("Converting pose encoding to extrinsic and intrinsic matrices...") extrinsic, intrinsic = pose_encoding_to_extri_intri(predictions["pose_enc"], images_tensor.shape[-2:]) predictions["extrinsic"] = extrinsic predictions["intrinsic"] = intrinsic for key in predictions.keys(): if isinstance(predictions[key], torch.Tensor): predictions[key] = predictions[key].cpu().numpy().squeeze(0) torch.cuda.empty_cache() return predictions def run_model(target_dir) -> dict: """ CPU-GPU hybrid: Handle CPU-intensive file I/O and call GPU function for inference. """ print(f"Processing images from {target_dir}") image_names = glob.glob(os.path.join(target_dir, "images", "*")) image_names = sorted(image_names) print(f"Found {len(image_names)} images") if len(image_names) == 0: raise ValueError("No images found. Check your upload.") images = load_and_preprocess_images(image_names) print(f"Preprocessed images shape: {images.shape}") predictions = _gpu_run_vggt_inference(images) print("Computing world points from depth map...") depth_map = predictions["depth"] # (S, H, W, 1) world_points = unproject_depth_map_to_point_map(depth_map, predictions["extrinsic"], predictions["intrinsic"]) predictions["world_points_from_depth"] = world_points return predictions def _prepare_upload_dir(input_file, input_video, frame_slider): """ CPU-only: Create target_dir, extract video frames or load PLY. No GPU usage. Returns (target_dir, image_paths). """ start_time = time.time() gc.collect() timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") target_dir = f"demo_output/inputs_{timestamp}" target_dir_images = os.path.join(target_dir, "images") target_dir_pcds = os.path.join(target_dir, "pcds") if os.path.exists(target_dir): shutil.rmtree(target_dir) os.makedirs(target_dir) os.makedirs(target_dir_images) os.makedirs(target_dir_pcds) image_paths = None if input_video is not None: print("processing video") if isinstance(input_video, dict) and "name" in input_video: video_path = input_video["name"] else: video_path = input_video vs = cv2.VideoCapture(video_path) fps = vs.get(cv2.CAP_PROP_FPS) frame_interval = int(fps * frame_slider) count = 0 video_frame_num = 0 image_paths = [] while True: gotit, frame = vs.read() if not gotit: break count += 1 if count % frame_interval == 0: image_path = os.path.join(target_dir_images, f"{video_frame_num:06}.png") cv2.imwrite(image_path, frame) image_paths.append(image_path) video_frame_num += 1 image_paths = sorted(image_paths) if input_file is not None: print("processing ply") original_points, original_colors, original_normals = _read_ply(input_file.name) if original_normals is None: original_normals = _estimate_normals_knn(original_points, k=30) scene_3d = trimesh.Scene() point_cloud_data = trimesh.PointCloud(vertices=original_points, colors=original_colors, vertex_normals=original_normals) scene_3d.add_geometry(point_cloud_data) original_temp = os.path.join(target_dir_pcds, "original.glb") scene_3d.export(file_obj=original_temp) np.save(os.path.join(target_dir_pcds, "points.npy"), original_points) np.save(os.path.join(target_dir_pcds, "colors.npy"), original_colors) np.save(os.path.join(target_dir_pcds, "normals.npy"), original_normals) end_time = time.time() print(f"Files prepared in {target_dir}; took {end_time - start_time:.3f} seconds") return target_dir, image_paths def handle_uploads(input_file, input_video, conf_thres, frame_slider, prediction_mode): """ Full pipeline: prepare files + run VGGT reconstruction (GPU) for video. Called by reconstruction_btn.click. """ start_time = time.time() target_dir, image_paths = _prepare_upload_dir(input_file, input_video, frame_slider) target_dir_pcds = os.path.join(target_dir, "pcds") if input_video is not None: original_points, original_colors, original_normals = parse_frames(target_dir, conf_thres, prediction_mode) scene_3d = trimesh.Scene() point_cloud_data = trimesh.PointCloud(vertices=original_points, colors=original_colors, vertex_normals=original_normals) scene_3d.add_geometry(point_cloud_data) original_temp = os.path.join(target_dir_pcds, "original.glb") scene_3d.export(file_obj=original_temp) np.save(os.path.join(target_dir_pcds, "points.npy"), original_points) np.save(os.path.join(target_dir_pcds, "colors.npy"), original_colors) np.save(os.path.join(target_dir_pcds, "normals.npy"), original_normals) else: original_temp = os.path.join(target_dir_pcds, "original.glb") end_time = time.time() return target_dir, image_paths, original_temp, end_time - start_time def update_gallery_on_upload(input_file, input_video, conf_thres, frame_slider, prediction_mode): """ Full reconstruction: called by reconstruction_btn.click. Includes GPU-based VGGT inference for video. """ if not input_video and not input_file: return None, None, None, None target_dir, image_paths, original_view, reconstruction_time = handle_uploads(input_file, input_video, conf_thres, frame_slider, prediction_mode) if input_file is not None: return original_view, target_dir, [], f"Upload and preprocess complete with {reconstruction_time:.3f} sec. Click \"PCA Generate\" to begin PCA processing." if input_video is not None: return original_view, target_dir, image_paths, f"Upload and preprocess complete with {reconstruction_time:.3f} sec. Click \"PCA Generate\" to begin PCA processing." def update_gallery_on_upload_cpu(input_file, input_video, _conf_thres, frame_slider, _prediction_mode): """ CPU-only upload: called by input_file.change / input_video.change. Only extracts frames and loads PLY, no GPU inference. """ if not input_video and not input_file: return None, None, None, None target_dir, image_paths = _prepare_upload_dir(input_file, input_video, frame_slider) target_dir_pcds = os.path.join(target_dir, "pcds") original_temp = os.path.join(target_dir_pcds, "original.glb") if input_file is not None: return original_temp, target_dir, [], f"Upload complete. Click \"PCA Generate\" to proceed." if input_video is not None: return None, target_dir, image_paths, f"Video frames extracted. Click \"Video Reconstruct\" to run 3D reconstruction." def clear_fields(): """ Clears the 3D viewer, the stored target_dir, and empties the gallery. """ return None def PCAing_log(is_example, log_output): """ Display a quick log message while waiting. """ if is_example: return log_output return "Loading for Doing PCA..." def reset_log(): """ Reset a quick log message. """ return "A new point cloud file or video is uploading and preprocessing..." def parse_frames( target_dir, conf_thres=3.0, prediction_mode="Pointmap Regression", ): """ Perform reconstruction using the already-created target_dir/images. """ if not os.path.isdir(target_dir) or target_dir == "None": return None, "No valid target directory found. Please upload first.", None, None start_time = time.time() gc.collect() # Prepare frame_filter dropdown target_dir_images = os.path.join(target_dir, "images") target_dir_pcds = os.path.join(target_dir, "pcds") all_files = sorted(os.listdir(target_dir_images)) if os.path.isdir(target_dir_images) else [] all_files = [f"{i}: {filename}" for i, filename in enumerate(all_files)] frame_filter_choices = ["All"] + all_files print("Running run_model...") predictions = run_model(target_dir) # Save predictions prediction_save_path = os.path.join(target_dir, "predictions.npz") np.savez(prediction_save_path, **predictions) # Convert pose encoding to extrinsic and intrinsic matrices images = predictions["images"] Ts, Ks = predictions["extrinsic"],predictions["intrinsic"] Ts = pad_0001(Ts) Ts_inv = np.linalg.inv(Ts) Cs = np.array([T_to_C(T) for T in Ts]) # (n, 3) # [1, 8, 294, 518, 3] world_points = predictions["world_points"] # Compute view direction for each pixel # (b n h w c) - (n, 3) view_dirs = world_points - rearrange(Cs, "n c -> n 1 1 c") view_dirs = rearrange(view_dirs, "n h w c -> (n h w) c") view_dirs = view_dirs / np.linalg.norm(view_dirs, axis=-1, keepdims=True) # Extract points and colors # [1, 8, 3, 294, 518] img_num = world_points.shape[1] images = predictions["images"] points = rearrange(world_points, "n h w c -> (n h w) c") colors = rearrange(images, "n c h w -> (n h w) c") if prediction_mode=="Pointmap Branch": world_points_conf = predictions["world_points_conf"] conf = world_points_conf.reshape(-1) points,Ts_inv,_ = Coord2zup(points, Ts_inv) scale = 3 / (points[:, 2].max() - points[:, 2].min()) points *= scale Ts_inv[:, :3, 3] *= scale # Create a point cloud pcd = PointCloud(points=points, colors=colors) pcd.estimate_normals(k=30) try: pcd, inliers, rotation_matrix, offset = extract_and_align_ground_plane(pcd) except Exception as e: print(f"cannot find ground, err:{e}") # Filp normals such that normals always point to camera # Compute the dot product between the normal and the view direction # If the dot product is less than 0, flip the normal normals = pcd.normals view_dirs = np.asarray(view_dirs) dot_product = np.sum(normals * view_dirs, axis=-1) flip_mask = dot_product > 0 normals[flip_mask] = -normals[flip_mask] # Normalize normals a nd m normals = normals / np.linalg.norm(normals, axis=-1, keepdims=True) pcd.normals = normals if conf_thres == 0.0: conf_threshold = 0.0 else: conf_threshold = np.percentile(conf, conf_thres) conf_mask = (conf >= conf_threshold) & (conf > 1e-5) points = points[conf_mask] colors = colors[conf_mask] normals = normals[conf_mask] elif prediction_mode=="Depthmap Branch": # Backproject per-frame depth maps into a fused world point cloud. # (n, h, w, 3) im_colors = rearrange(images, "n c h w -> (n) h w c") # (n, h, w) im_dists = world_points - rearrange(Cs, "n c -> n 1 1 c") im_dists = np.linalg.norm(im_dists, axis=-1, keepdims=False) # Convert distance to depth im_depths = [] # (n, h, w) for im_dist, K in zip(im_dists, Ks): im_depths.append(im_distance_to_im_depth(im_dist, K)) im_depths = np.stack(im_depths, axis=0) points = [] for K, T, im_depth in zip(Ks, Ts, im_depths): points.append(im_depth_to_point_cloud(im_depth, K, T)) points = np.vstack(points) colors = im_colors.reshape(-1, 3) world_points_conf = predictions["depth_conf"] conf = world_points_conf.reshape(-1) if conf_thres == 0.0: conf_threshold = 0.0 else: conf_threshold = np.percentile(conf, conf_thres) conf_mask = (conf >= conf_threshold) & (conf > 1e-5) points = points[conf_mask] colors = colors[conf_mask] points, Ts_inv, _ = Coord2zup(points, Ts_inv) scale_factor = 3. / (np.max(points[:, 2]) - np.min(points[:, 2])) points *= scale_factor Ts_inv[:, :3, 3] *= scale_factor pcd = PointCloud(points=points, colors=colors) pcd.estimate_normals(k=30) try: pcd, inliers, rotation_matrix, offset = extract_and_align_ground_plane(pcd) except Exception as e: print(f"cannot find ground, err:{e}") original_points = pcd.points original_colors = pcd.colors original_normals = pcd.normals # Cleanup del predictions gc.collect() end_time = time.time() print(f"Total time: {end_time - start_time:.2f} seconds") return original_points, original_colors, original_normals def extract_and_align_ground_plane(pcd, height_percentile=20, ransac_distance_threshold=0.01, ransac_n=3, ransac_iterations=1000, max_angle_degree=40, max_trials=6): points = pcd.points z_vals = points[:, 2] z_thresh = np.percentile(z_vals, height_percentile) low_indices = np.where(z_vals <= z_thresh)[0] remaining_indices = low_indices.copy() for trial in range(max_trials): if len(remaining_indices) < ransac_n: raise ValueError("Not enough points left to fit a plane.") low_points = points[remaining_indices] plane_model, inliers = _segment_plane_ransac( low_points, distance_threshold=ransac_distance_threshold, ransac_n=ransac_n, num_iterations=ransac_iterations, ) a, b, c, d = plane_model normal = np.array([a, b, c]) normal /= np.linalg.norm(normal) angle = np.arccos(np.clip(np.dot(normal, [0, 0, 1]), -1.0, 1.0)) * 180 / np.pi if angle <= max_angle_degree: inliers_global = remaining_indices[inliers] target = np.array([0, 0, 1]) axis = np.cross(normal, target) axis_norm = np.linalg.norm(axis) if axis_norm < 1e-6: rotation_matrix = np.eye(3) else: axis /= axis_norm rot_angle = np.arccos(np.clip(np.dot(normal, target), -1.0, 1.0)) rotation = R.from_rotvec(axis * rot_angle) rotation_matrix = rotation.as_matrix() rotated_points = points @ rotation_matrix.T ground_points_z = rotated_points[inliers_global, 2] offset = np.mean(ground_points_z) rotated_points[:, 2] -= offset aligned_pcd = PointCloud(points=rotated_points) if pcd.has_colors(): aligned_pcd.colors = pcd.colors if pcd.has_normals(): rotated_normals = pcd.normals @ rotation_matrix.T aligned_pcd.normals = rotated_normals return aligned_pcd, inliers_global, rotation_matrix, offset else: rejected_indices = remaining_indices[inliers] remaining_indices = np.setdiff1d(remaining_indices, rejected_indices) raise ValueError("Failed to find a valid ground plane within max trials.") def rotx(x, theta=90): """ Rotate x by theta degrees around the x-axis """ theta = np.deg2rad(theta) rot_matrix = np.array( [ [1, 0, 0, 0], [0, np.cos(theta), -np.sin(theta), 0], [0, np.sin(theta), np.cos(theta), 0], [0, 0, 0, 1], ] ) return rot_matrix@ x def Coord2zup(points, extrinsics, normals = None): """ Convert the dust3r coordinate system to the z-up coordinate system """ points = np.concatenate([points, np.ones([points.shape[0], 1])], axis=1).T points = rotx(points, -90)[:3].T if normals is not None: normals = np.concatenate([normals, np.ones([normals.shape[0], 1])], axis=1).T normals = rotx(normals, -90)[:3].T normals = normals / np.linalg.norm(normals, axis=1, keepdims=True) t = np.min(points,axis=0) points -= t extrinsics = rotx(extrinsics, -90) extrinsics[:, :3, 3] -= t.T return points, extrinsics, normals def get_pca_color(feat, start = 0, brightness=1.25, center=True): u, s, v = torch.pca_lowrank(feat, center=center, q=3*(start+1), niter=5) projection = feat @ v projection = projection[:, 3*start:3*(start+1)] * 0.6 + projection[:, 3*start:3*(start+1)] * 0.4 min_val = projection.min(dim=-2, keepdim=True)[0] max_val = projection.max(dim=-2, keepdim=True)[0] div = torch.clamp(max_val - min_val, min=1e-6) color = (projection - min_val) / div * brightness color = color.clamp(0.0, 1.0) return color @spaces.GPU def _gpu_concerto_forward_pca(point, concerto_model_, pca_slider, bright_slider): """ GPU-only function: Run Concerto/Sonata model forward pass and PCA. """ device = "cuda" if torch.cuda.is_available() else "cpu" for key in point.keys(): if isinstance(point[key], torch.Tensor): point[key] = point[key].to(device, non_blocking=True) concerto_model_ = concerto_model_.to(device) concerto_model_.eval() with torch.inference_mode(): concerto_start_time = time.time() with torch.inference_mode(False): point = concerto_model_(point) concerto_end_time = time.time() # upcast point feature for _ in range(2): assert "pooling_parent" in point.keys() assert "pooling_inverse" in point.keys() parent = point.pop("pooling_parent") inverse = point.pop("pooling_inverse") parent.feat = torch.cat([parent.feat, point.feat[inverse]], dim=-1) point = parent while "pooling_parent" in point.keys(): assert "pooling_inverse" in point.keys() parent = point.pop("pooling_parent") inverse = point.pop("pooling_inverse") parent.feat = point.feat[inverse] point = parent pca_start_time = time.time() pca_color = get_pca_color(point.feat, start=pca_slider, brightness=bright_slider, center=True) pca_end_time = time.time() original_pca_color = pca_color[point.inverse] processed_colors = original_pca_color.cpu().detach().numpy() point_feat = point.feat.cpu().detach().numpy() point_inverse = point.inverse.cpu().detach().numpy() concerto_time = concerto_end_time - concerto_start_time pca_time = pca_end_time - pca_start_time torch.cuda.empty_cache() return processed_colors, point_feat, point_inverse, concerto_time, pca_time def Concerto_process(target_dir, original_points, original_colors, original_normals, slider_value, bright_value, model_type): target_dir_pcds = os.path.join(target_dir, "pcds") point = {"coord": original_points, "color": original_colors, "normal": original_normals} original_coord = point["coord"].copy() point = transform(point) # Select model based on type if model_type == "Concerto": selected_model = concerto_model elif model_type == "Sonata": selected_model = sonata_model else: selected_model = concerto_model # GPU: Run model forward + PCA processed_colors, point_feat, point_inverse, concerto_time, pca_time = _gpu_concerto_forward_pca( point, selected_model, slider_value, bright_value ) # CPU: Save features np.save(os.path.join(target_dir_pcds, "feat.npy"), point_feat) np.save(os.path.join(target_dir_pcds, "inverse.npy"), point_inverse) return original_coord, processed_colors, concerto_time, pca_time def gradio_demo(target_dir,pca_slider,bright_slider, model_type, if_color=True, if_normal=True): if target_dir is None or target_dir == "None": return None, "No point cloud available. Please upload data first." target_dir_pcds = os.path.join(target_dir, "pcds") if not os.path.isfile(os.path.join(target_dir_pcds,"points.npy")): return None, "No point cloud available. Please upload data first." original_points = np.load(os.path.join(target_dir_pcds,"points.npy")) if if_color: original_colors = np.load(os.path.join(target_dir_pcds,"colors.npy")) else: original_colors = np.zeros_like(original_points) if if_normal: original_normals = np.load(os.path.join(target_dir_pcds,"normals.npy")) else: original_normals = np.zeros_like(original_points) processed_temp = (os.path.join(target_dir_pcds,"processed.glb")) processed_points, processed_colors, concerto_time, pca_time = Concerto_process(target_dir,original_points, original_colors,original_normals, pca_slider, bright_slider, model_type) feat_3d = trimesh.Scene() feat_data = trimesh.PointCloud(vertices=processed_points, colors=processed_colors, vertex_normals=original_normals) feat_3d.add_geometry(feat_data) feat_3d.export(processed_temp) return processed_temp, f"Feature visualization process finished with {concerto_time:.3f} seconds using Concerto inference and {pca_time:.3f} seconds using PCA. Updating visualization." @spaces.GPU def _gpu_pca_slider_compute(feat_array, inverse_array, pca_slider, bright_slider): """ GPU-only function: Compute PCA colors for slider updates. """ device = "cuda" if torch.cuda.is_available() else "cpu" feat_tensor = torch.tensor(feat_array, device=device) inverse_tensor = torch.tensor(inverse_array, device=device) pca_start_time = time.time() pca_colors = get_pca_color(feat_tensor, start=pca_slider, brightness=bright_slider, center=True) processed_colors = pca_colors[inverse_tensor].cpu().detach().numpy() pca_end_time = time.time() return processed_colors, (pca_end_time - pca_start_time) def concerto_slider_update(target_dir,pca_slider,bright_slider,is_example,log_output): if is_example == "True": return None, log_output else: target_dir_pcds = os.path.join(target_dir, "pcds") if os.path.isfile(os.path.join(target_dir_pcds,"feat.npy")): # CPU: Load data from disk feat = np.load(os.path.join(target_dir_pcds,"feat.npy")) inverse = np.load(os.path.join(target_dir_pcds,"inverse.npy")) # GPU: Compute PCA colors processed_colors, pca_time = _gpu_pca_slider_compute(feat, inverse, pca_slider, bright_slider) # CPU: Build mesh processed_points = np.load(os.path.join(target_dir_pcds,"points.npy")) processed_normals = np.load(os.path.join(target_dir_pcds,"normals.npy")) processed_temp = (os.path.join(target_dir_pcds,"processed.glb")) feat_3d = trimesh.Scene() feat_data = trimesh.PointCloud(vertices=processed_points, colors=processed_colors, vertex_normals=processed_normals) feat_3d.add_geometry(feat_data) feat_3d.export(processed_temp) log_output = f"Feature visualization process finished with {pca_time:.3f} seconds using PCA. Updating visualization." else: processed_temp = None log_output = "No representations saved, please click PCA generate first." return processed_temp, log_output BASE_URL = "https://huggingface.co/datasets/pointcept-bot/concerto_huggingface_demo/resolve/main/" def get_url(path): return f"{BASE_URL}{path}" examples_video = [ [get_url("video/re10k_1.mp4"), 10.0, 1, "Depthmap Branch", 2, 1.2, "True"], [get_url("video/re10k_2.mp4"), 30.0, 1, "Depthmap Branch", 1, 1.2, "True"], [get_url("video/re10k_3.mp4"), 10.0, 1, "Depthmap Branch", 1, 1.2, "True"], [get_url("video/re10k_4.mp4"), 10.0, 1, "Depthmap Branch", 1, 1.0, "True"], ] examples_pcd = [ [get_url("pcd/scannet_0024.png"),get_url("pcd/scannet_0024.ply"),2,1.2, "True"], [get_url("pcd/scannet_0603.png"),get_url("pcd/scannet_0603.ply"),0,1.2, "True"], [get_url("pcd/hm3d_00113_3goH1WRaCYC.png"),get_url("pcd/hm3d_00113_3goH1WRaCYC.ply"),0,1.2, "True"], [get_url("pcd/s3dis_Area2_auditorium1.png"),get_url("pcd/s3dis_Area2_auditorium1.ply"),0,1.2, "True"], ] with gr.Blocks( css=""" .custom-log * { font-style: italic; font-size: 22px !important; background-image: linear-gradient(120deg, #0ea5e9 0%, #6ee7b7 60%, #34d399 100%); -webkit-background-clip: text; background-clip: text; font-weight: bold !important; color: transparent !important; text-align: center !important; width: 800px; height: 100px; } .example-log * { font-style: italic; font-size: 16px !important; background-image: linear-gradient(120deg, #0ea5e9 0%, #6ee7b7 60%, #34d399 100%); -webkit-background-clip: text; background-clip: text; color: transparent !important; } .common-markdown * { font-size: 22px !important; -webkit-background-clip: text; background-clip: text; font-weight: bold !important; color: #0ea5e9 !important; text-align: center !important; } #big-box { border: 3px solid #00bcd4; padding: 20px; background-color: transparent; border-radius: 15px; } #my_radio .wrap { display: flex; flex-wrap: nowrap; justify-content: center; align-items: center; } #my_radio .wrap label { display: flex; width: 50%; justify-content: center; align-items: center; margin: 0; padding: 10px 0; box-sizing: border-box; } """, ) as demo: gr.HTML( """

Concerto: Joint 2D-3D Self-Supervised Learning for Emergent Spatial Representations

    Getting Started:(Click to expand)
  1. Before Start:Due to space limitations, you may encounter errors. If so, please deploy this demo locally.
  2. Upload Your Data: Use the "Upload Video" or "Upload Point Cloud" blocks on the left to provide your input. If you upload a video, it will be automatically split into individual frames with the specified frame gap by VGGT.
  3. [Optional] Adjust Video-Lifted Point Cloud: Before reconstructing the video, you can fine-tune the VGGT lifting process using the options below
    (Click to expand)
    • Frame Gap / N Sec: Adjust the frame interval.
    • Confidence Threshold: Adjust the point filtering based on confidence levels.
    • Select Prediction Mode: Choose between "Depthmap Branch" and "Pointmap Branch."
  4. PCA Generation: After reconstruction, click the "PCA Generate" button to start the representation extraction and PCA process.
  5. Clear: Click the "Clear" button to reset all content in the blocks.
  6. Point Cloud Preview: Your uploaded video or point cloud will be displayed in this block.
  7. PCA Result: The PCA point cloud will appear here. You can rotate, drag, and zoom to explore the model, and download the GLB file.
  8. [Optional] Adjust the Point Cloud Input (pre-release feature of the next work): use the checkbox "Input with Point Cloud Color" and "Input with Point Cloud Normal".
  9. [Optional] Adjust PCA Visualization: Fine-tune the PCA visualization using the options below
    (Click to expand)
    • Model Type: Choose the model from Concerto and Sonata.
    • PCA Start Dimension: PCA reduces high-dimensional representations into 3D vectors. Adjust the PCA start dimension to change the range of the visualization. Increasing this value can help you see PCA visualization with less variance when the initial PCA dimension shows less diversity.
    • PCA Brightness: Adjust the brightness of the PCA visualization results.
    • Notice: As a linear dimension reduction method, PCA has its limitation. Sometimes, the visualization cannot fully exhibit the quality of representations.
""" ) _ = gr.Textbox(label="_", visible=False, value="False") is_example = gr.Textbox(label="is_example", visible=False, value="False") target_dir = gr.Textbox(label="Target Dir", visible=False, value="None") preview_imgs = gr.Image(type="filepath",label="Preview Imgs", visible=False, value="None") with gr.Row(): with gr.Column(scale=1,elem_id="big-box"): input_file = gr.File(label="Upload Point Cloud", file_types=[".ply"]) input_video = gr.Video(label="Upload Video", interactive=True) image_gallery = gr.Gallery( label="Preview", columns=4, height="300px", object_fit="contain", preview=True, ) frame_slider = gr.Slider(minimum=0.1, maximum=10, value=1, step=0.1, label="1 Frame/ N Sec", interactive=True) conf_thres = gr.Slider(minimum=0, maximum=100, value=10, step=0.1, label="Confidence", interactive=True) prediction_mode = gr.Radio( ["Depthmap Branch", "Pointmap Branch"], label="Select a Prediction Mode", value="Depthmap Branch", scale=1, elem_id="my_radio", ) reconstruction_btn = gr.Button("Video Reconstruct") with gr.Column(scale=2): log_output = gr.Markdown( "Please upload a video or point cloud ply file, then click \"PCA Generate\".", elem_classes=["custom-log"] ) original_view = gr.Model3D(height=520, zoom_speed=0.5, pan_speed=0.5, label="Point Cloud Preview", camera_position = (90,None,None)) processed_view = gr.Model3D(height=520, zoom_speed=0.5, pan_speed=0.5, label="PCA Result", camera_position = (90,None,None)) with gr.Row(): if_color = gr.Checkbox(label="Input with Point Cloud Color", value=True) if_normal = gr.Checkbox(label="Input with Point Cloud Normal", value=True) model_type = gr.Radio( ["Concerto", "Sonata"], label="Select a Model Type", value="Concerto", scale=1, elem_id="my_radio", ) pca_slider = gr.Slider(minimum=0, maximum=5, value=0, step=1, label="PCA Start Dimension", interactive=True) bright_slider = gr.Slider(minimum=0.5, maximum=1.5, value=1.2, step=0.05, label="PCA Brightness", interactive=True) with gr.Row(): submit_btn = gr.Button("PCA Generate") clear_btn = gr.ClearButton( [input_video, input_file, original_view, processed_view, log_output, target_dir, image_gallery], scale=1, elem_id="my_clear", ) gr.Markdown("Click any row to load an example.", elem_classes=["example-log"]) with gr.Row(): def example_video_updated( inputs, conf_thres, frame_slider, prediction_mode, pca_slider, bright_slider, is_example, ): return inputs,conf_thres,frame_slider,prediction_mode,pca_slider,bright_slider,is_example gr.Examples( examples=examples_video, inputs=[ input_video, conf_thres, frame_slider, prediction_mode, pca_slider, bright_slider, is_example, ], outputs=[ input_video, conf_thres, frame_slider, prediction_mode, pca_slider, bright_slider, is_example, ], label = "Video Examples", fn=example_video_updated, cache_examples=False, examples_per_page=50, # examples_per_page=2 ) with gr.Row(): def example_file_updated( preview_imgs, inputs, pca_slider, bright_slider, is_example, ): return inputs,pca_slider,bright_slider,is_example gr.Examples( examples=examples_pcd, inputs=[ preview_imgs, input_file, pca_slider, bright_slider, is_example, ], outputs=[ input_file, pca_slider, bright_slider, is_example, ], label = "Point Cloud Examples", fn=example_file_updated, cache_examples=False, examples_per_page=50, # examples_per_page=2 ) reconstruction_btn.click( fn = update_gallery_on_upload, inputs = [input_file,input_video,conf_thres,frame_slider,prediction_mode], outputs = [original_view, target_dir, image_gallery, log_output] ) submit_btn.click(fn=clear_fields, inputs=[], outputs=[processed_view]).then( fn=PCAing_log, inputs=[is_example, log_output], outputs=[log_output] ).then( fn=gradio_demo, inputs=[target_dir,pca_slider,bright_slider, model_type, if_color, if_normal], outputs=[processed_view,log_output], ).then( fn=lambda: "False", inputs=[], outputs=[is_example] # set is_example to "False" ) pca_slider.release(fn=clear_fields, inputs=[], outputs=[processed_view]).then( fn=PCAing_log, inputs=[is_example, log_output], outputs=[log_output] ).then( fn=concerto_slider_update, inputs=[target_dir,pca_slider,bright_slider,is_example,log_output], outputs=[processed_view, log_output], ).then( fn=lambda: "False", inputs=[], outputs=[is_example] # set is_example to "False" ) bright_slider.release(fn=clear_fields, inputs=[], outputs=[processed_view]).then( fn=PCAing_log, inputs=[is_example, log_output], outputs=[log_output] ).then( fn=concerto_slider_update, inputs=[target_dir,pca_slider,bright_slider,is_example,log_output], outputs=[processed_view, log_output], ).then( fn=lambda: "False", inputs=[], outputs=[is_example] # set is_example to "False" ) model_type.change(fn=clear_fields, inputs=[], outputs=[processed_view]).then( fn=PCAing_log, inputs=[is_example, log_output], outputs=[log_output] ).then( fn=gradio_demo, inputs=[target_dir,pca_slider,bright_slider, model_type, if_color, if_normal], outputs=[processed_view,log_output], ).then( fn=lambda: "False", inputs=[], outputs=[is_example] # set is_example to "False" ) input_file.change(fn=reset_log, inputs=[], outputs=[log_output]).then( fn=update_gallery_on_upload_cpu, inputs=[input_file,input_video, conf_thres,frame_slider,prediction_mode], outputs=[original_view, target_dir, _, log_output], ) input_video.change(fn=reset_log, inputs=[], outputs=[log_output]).then( fn=update_gallery_on_upload_cpu, inputs=[input_file,input_video, conf_thres,frame_slider,prediction_mode], outputs=[original_view, target_dir, image_gallery, log_output], ) if __name__ == "__main__": demo.queue(max_size=20).launch(show_error=True, share=True)