app.py

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(
    """
    <h1>Concerto: Joint 2D-3D Self-Supervised Learning for Emergent Spatial Representations</h1>
    <div style="font-size: 16px; line-height: 1.5;">
        <ol>
            <details style="display:inline;">
                <summary style="display:inline;"><h3>Getting Started:(<strong>Click to expand</strong>)</h3></summary>
            <li><strong>Before Start:</strong>Due to space limitations, you may encounter errors. If so, please deploy this demo locally.</li>
            <li><strong>Upload Your Data:</strong> 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.</li>
            <li>
                <strong>[Optional] Adjust Video-Lifted Point Cloud:</strong>
                Before reconstructing the video, you can fine-tune the VGGT lifting process using the options below
                <details style="display:inline;">
                    <summary style="display:inline;">(<strong>Click to expand</strong>)</summary>
                    <ul>
                        <li><em>Frame Gap / N Sec:</em> Adjust the frame interval.</li>
                        <li><em>Confidence Threshold:</em> Adjust the point filtering based on confidence levels.</li>
                        <li><em>Select Prediction Mode:</em> Choose between "Depthmap Branch" and "Pointmap Branch."</li>
                    </ul>
                </details>
            </li>
            <li><strong>PCA Generation:</strong> After reconstruction, click the "PCA Generate" button to start the representation extraction and PCA process.</li>
            <li><strong>Clear:</strong> Click the "Clear" button to reset all content in the blocks.</li>
            <li><strong>Point Cloud Preview:</strong> Your uploaded video or point cloud will be displayed in this block.</li>
            <li><strong>PCA Result:</strong> The PCA point cloud will appear here. You can rotate, drag, and zoom to explore the model, and download the GLB file.</li>
            <li>
                <strong>[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".
            </li>
            <li>
                <strong>[Optional] Adjust PCA Visualization:</strong>
                Fine-tune the PCA visualization using the options below
                <details style="display:inline;">
                    <summary style="display:inline;">(<strong>Click to expand</strong>)</summary>
                    <ul>
                        <li><em>Model Type:</em> Choose the model from Concerto and Sonata.</li>
                        <li><em>PCA Start Dimension:</em> 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.</li>
                        <li><em>PCA Brightness:</em> Adjust the brightness of the PCA visualization results.</li>
                        <li><em>Notice:</em> As a linear dimension reduction method, PCA has its limitation. Sometimes, the visualization cannot fully exhibit the quality of representations.</li>
                    </ul>
                </details>
            </li>
            </details>
        </ol>
    </div>

    """
    )
    _ = 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)