app.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os, glob, re, torch, numpy as np
from typing import List, Tuple
from safetensors.torch import load_file
import gradio as gr
import plotly.graph_objects as go
import torch.nn as nn
from huggingface_hub import hf_hub_download, list_repo_files

# =========================
# Config
# =========================
LOCAL_FILE_DEFAULT  = "assets/scene0073_00.safetensors"  # local safetensors file
PM_KEY_DEFAULT      = "point_map"
TOPK_VIEWS_DEFAULT  = 3

VOXEL_SIZE          = 0.02
DOWNSAMPLE_N_MAX    = 600_000
POINT_SIZE_MIN      = 1.2
POINT_SIZE_MAX      = 2.0

CLR_RED             = "rgba(230,40,40,0.98)"
BG_COLOR            = "#f7f9fb"
GRID_COLOR          = "#e6ecf2"
BOX_COLOR           = "rgba(80,80,80,0.6)"

TOP_VIEW_IMAGE_PATH = "assets/scene0073_00.png"

DEFAULT_CAM = dict(
    eye=dict(x=1.35, y=1.35, z=0.95),
    up=dict(x=0, y=0, z=1),
    center=dict(x=0, y=0, z=0),
    projection=dict(type="perspective"),
)

def _merge_safetensors_dicts(paths: List[str]):
    merged = {}
    for p in paths:
        sd = load_file(p, device="cpu")
        merged.update(sd)
    return merged

def _local_all_under(path: str) -> List[str]:
    out = []
    if os.path.isfile(path):
        return [path]
    for root, _, files in os.walk(path):
        for f in files:
            out.append(os.path.join(root, f))
    return sorted(out)

# =========================
# Load pretrained (your existing loader)
# =========================
def load_pretrain(
    model: torch.nn.Module,
    ckpt_path: str,                 # e.g. "assets/ckpt_100.pth" or "assets/model.safetensors"
    repo_id: str = "MatchLab/poma3d-demo",
    revision: str = "main",
    allow_local_fallback: bool = True,
):
    if allow_local_fallback and (os.path.isfile(ckpt_path) or os.path.isdir(ckpt_path)):
        print(f"📂 Using local checkpoint(s): {ckpt_path}")
        local_files = _local_all_under(ckpt_path)
    else:
        # 2) REMOTE: resolve file list from Space
        print(f"📦 Resolving from HF Space: {repo_id}/{ckpt_path} (rev={revision})")
        files = list_repo_files(repo_id=repo_id, repo_type='model', revision=revision)

        # Exact file hit?
        if ckpt_path in files:
            to_fetch = [ckpt_path]
        else:
            # Treat ckpt_path as a folder prefix (ensure trailing slash for matching)
            prefix = ckpt_path if ckpt_path.endswith("/") else ckpt_path + "/"
            to_fetch = [f for f in files if f.startswith(prefix)]
            if not to_fetch:
                preview = "\n".join(files[:100])
                raise FileNotFoundError(
                    f"'{ckpt_path}' not found in Space '{repo_id}' (rev='{revision}').\n"
                    f"Files present (first 100):\n{preview}"
                )

        # Download all matching files locally
        local_files = []
        for rel in to_fetch:
            lp = hf_hub_download(repo_id=repo_id, filename=rel, repo_type='model', revision=revision)
            local_files.append(lp)
        local_files.sort()

    # Filter by types we know how to load
    safes = [p for p in local_files if p.endswith(".safetensors")]
    pths  = [p for p in local_files if re.search(r"\.(?:pth|pt)$", p)]

    if safes:
        print(f"🧩 Found {len(safes)} .safetensors shard(s); merging…")
        state = _merge_safetensors_dicts(safes)
    elif pths:
        # pick the largest .pth/.pt to avoid optimizer/state variants
        pths_sorted = sorted(pths, key=lambda p: os.path.getsize(p), reverse=True)
        pick = pths_sorted[0]
        print(f"🧩 Using .pth/.pt: {os.path.basename(pick)} (largest of {len(pths)} candidates)")
        state = torch.load(pick, map_location="cpu")
        # strip common prefixes
        if isinstance(state, dict) and any(k.startswith(("model.", "target_model.")) for k in state.keys()):
            state = { (k.split(".", 1)[1] if k.startswith(("model.", "target_model.")) else k): v
                      for k, v in state.items() }
        # nested 'state_dict'
        if isinstance(state, dict) and "state_dict" in state and isinstance(state["state_dict"], dict):
            state = state["state_dict"]
    else:
        raise FileNotFoundError(
            "No loadable checkpoint found. Expecting one or more of: "
            ".safetensors or .pth/.pt under the given path."
        )

    # Load into model
    result = model.load_state_dict(state, strict=False)

    # Report
    weight_keys = set(state.keys()) if isinstance(state, dict) else set()
    model_keys  = set(model.state_dict().keys())
    loaded_keys = model_keys.intersection(weight_keys)
    print("✅ Weights loaded")
    print(f"   • Loaded keys: {len(loaded_keys)}")
    print(f"   • Missing keys: {len(result.missing_keys)}")
    print(f"   • Unexpected keys: {len(result.unexpected_keys)}")

    return result
# =========================
# Representation model (fg-clip-base + LoRA)
# =========================
def build_model(device: torch.device):
    from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
    from peft import LoraConfig, get_peft_model

    class RepModel(torch.nn.Module):
        def __init__(self, model_root="qihoo360/fg-clip-base"):
            super().__init__()
            lora_cfg = LoraConfig(
                r=32, lora_alpha=64,
                target_modules=["q_proj","k_proj","v_proj","fc1","fc2"],
                lora_dropout=0.05, bias="none",
                task_type="FEATURE_EXTRACTION"
            )
            cfg = AutoConfig.from_pretrained(model_root, trust_remote_code=True)
            base = AutoModelForCausalLM.from_config(cfg, trust_remote_code=True)
            self.target_model = get_peft_model(base, lora_cfg)
            self.tokenizer = AutoTokenizer.from_pretrained(model_root, trust_remote_code=True, use_fast=True)

        @torch.no_grad()
        def get_text_feature(self, texts, device):
            tok = self.tokenizer(texts, padding="max_length", truncation=True, max_length=248, return_tensors="pt").to(device)
            feats = self.target_model.get_text_features(tok["input_ids"], walk_short_pos=False)
            feats = torch.nn.functional.normalize(feats.float(), dim=-1)
            return feats

        @torch.no_grad()
        def get_image_feature(self, pm_batched):
            feats = self.target_model.get_image_features(pm_batched)
            feats = torch.nn.functional.normalize(feats.float(), dim=-1)
            return feats

    m = RepModel().to(device).eval()
    print("Using fg-clip-base RepModel.")
    return m

# =========================
# Data loading & helpers
# =========================
def load_scene_local(path: str, pm_key: str = PM_KEY_DEFAULT) -> torch.Tensor:
    if not os.path.exists(path):
        raise FileNotFoundError(f"Local file not found: {path}")
    sd = load_file(path)
    if pm_key not in sd:
        raise KeyError(f"Key '{pm_key}' not found in {list(sd.keys())}")
    pm = sd[pm_key]  # (V,H,W,3)
    if pm.dim() != 4 or pm.shape[-1] != 3:
        raise ValueError(f"Invalid shape {tuple(pm.shape)}, expected (V,H,W,3)")
    return pm.permute(0, 3, 1, 2).contiguous()  # -> (V,3,H,W)

def _xyz_to_numpy(xyz: torch.Tensor) -> np.ndarray:
    pts = xyz.permute(1, 2, 0).reshape(-1, 3).cpu().numpy().astype(np.float32)
    mask = np.isfinite(pts).all(axis=1)
    return pts[mask]

def stack_views(pm: torch.Tensor) -> Tuple[np.ndarray, np.ndarray]:
    pts_all, vid_all = [], []
    for v in range(pm.shape[0]):
        pts = _xyz_to_numpy(pm[v])
        if pts.size == 0: continue
        pts_all.append(pts)
        vid_all.append(np.full((pts.shape[0],), v, dtype=np.int32))
    pts_all = np.concatenate(pts_all, axis=0)
    vid_all = np.concatenate(vid_all, axis=0)
    return pts_all, vid_all

def voxel_downsample_with_ids(pts, vids, voxel: float):
    if pts.shape[0] == 0: return pts, vids
    grid = np.floor(pts / voxel).astype(np.int64)
    key = np.core.records.fromarrays(grid.T, names="x,y,z", formats="i8,i8,i8")
    _, uniq_idx = np.unique(key, return_index=True)
    return pts[uniq_idx], vids[uniq_idx]

def hard_cap(pts, vids, cap: int):
    N = pts.shape[0]
    if N <= cap: return pts, vids
    idx = np.random.choice(N, size=cap, replace=False)
    return pts[idx], vids[idx]

def adaptive_point_size(n: int) -> float:
    ps = 2.4 * (150_000 / max(n, 10)) ** 0.25
    return float(np.clip(ps, POINT_SIZE_MIN, POINT_SIZE_MAX))

def scene_bbox(pts: np.ndarray):
    mn, mx = pts.min(axis=0), pts.max(axis=0)
    x0,y0,z0 = mn; x1,y1,z1 = mx
    corners = np.array([
        [x0,y0,z0],[x1,y0,z0],[x1,y1,z0],[x0,y1,z0],
        [x0,y0,z1],[x1,y0,z1],[x1,y1,z1],[x0,y1,z1]
    ])
    edges = [(0,1),(1,2),(2,3),(3,0),(4,5),(5,6),(6,7),(7,4),(0,4),(1,5),(2,6),(3,7)]
    xs,ys,zs=[],[],[]
    for a,b in edges:
        xs += [corners[a,0], corners[b,0], None]
        ys += [corners[a,1], corners[b,1], None]
        zs += [corners[a,2], corners[b,2], None]
    return xs,ys,zs

@torch.no_grad()
def rank_views_for_text(model, text, pm, device, topk: int):
    img_feats = model.get_image_feature(pm.float().to(device))
    txt_feat  = model.get_text_feature([text], device=device)[0]
    sims = torch.matmul(img_feats, txt_feat)
    order = torch.argsort(sims, descending=True)[:max(1, int(topk))]
    return order.tolist()

# =========================
# Visualization
# =========================
def depth_values(pts: np.ndarray) -> np.ndarray:
    z = pts[:, 2]
    z_min, z_max = z.min(), z.max()
    return (z - z_min) / (z_max - z_min + 1e-9)

def base_figure_gray_depth(pts: np.ndarray, point_size: float, camera=DEFAULT_CAM) -> go.Figure:
    depth = depth_values(pts)
    fig = go.Figure(go.Scatter3d(
        x=pts[:,0], y=pts[:,1], z=pts[:,2],
        mode="markers",
        marker=dict(size=point_size, color=depth, colorscale="Greys", reversescale=True, opacity=0.50),
        hoverinfo="skip"
    ))
    bx,by,bz = scene_bbox(pts)
    fig.add_trace(go.Scatter3d(x=bx,y=by,z=bz,mode="lines",line=dict(color=BOX_COLOR,width=2),hoverinfo="skip"))
    fig.update_layout(scene=dict(aspectmode="data",camera=camera),
                      margin=dict(l=0,r=0,b=0,t=0),
                      paper_bgcolor=BG_COLOR,
                      showlegend=False)
    return fig

def highlight_views_3d(pts, view_ids, selected, point_size, camera=DEFAULT_CAM):
    depth = depth_values(pts)
    colors = np.stack([depth, depth, depth], axis=1)
    if selected:
        sel_mask = np.isin(view_ids, np.array(selected, dtype=np.int32))
        colors[sel_mask] = np.array([1, 0, 0])
    fig = go.Figure(go.Scatter3d(
        x=pts[:,0], y=pts[:,1], z=pts[:,2],
        mode="markers",
        marker=dict(size=point_size,
                    color=[f"rgb({int(r*255)},{int(g*255)},{int(b*255)})"
                           for r,g,b in colors],
                    opacity=0.98),
        hoverinfo="skip"
    ))
    bx,by,bz = scene_bbox(pts)
    fig.add_trace(go.Scatter3d(x=bx,y=by,z=bz,mode="lines",
                               line=dict(color=BOX_COLOR,width=2),hoverinfo="skip"))
    fig.update_layout(scene=dict(aspectmode="data",camera=camera),
                      margin=dict(l=0,r=0,b=0,t=0),
                      paper_bgcolor=BG_COLOR,
                      showlegend=False)
    return fig

# =========================
# App setup
# =========================
device = 'cpu'
model  = build_model(device)
load_pretrain(model, "assets/ckpt_100.pth")

with gr.Blocks(
    title="POMA-3D: Text-conditioned 3D Scene Visualization",
    css="#plot3d, #img_ref {height: 450px !important;}"
) as demo:
    gr.Markdown("### POMA-3D: The Point Map Way to 3D Scene Understanding - Embodied Localization Demo\n"
                "Enter agent's situation text and choose **Top-K**; the most relevant views will turn **red**.")

    with gr.Row():
        text_in = gr.Textbox(label="Text query", value="I am sleeping on the bed.", scale=4)
        topk_in = gr.Number(label="Top-K views", value=TOPK_VIEWS_DEFAULT, precision=0, minimum=1, maximum=12)
        submit_btn = gr.Button("Locate", variant="primary")

    with gr.Row(equal_height=True):
        with gr.Column(scale=1, min_width=500):
            plot3d = gr.Plot(label="3D Point Cloud (rotatable)", elem_id="plot3d")
        with gr.Column(scale=1, min_width=500):
            img_ref = gr.Image(label="Top-Down Reference View", value=TOP_VIEW_IMAGE_PATH, elem_id="img_ref")

    status = gr.Markdown()

    pm_state   = gr.State(None)
    pts_state  = gr.State(None)
    vids_state = gr.State(None)

    # Load scene automatically from LOCAL_FILE_DEFAULT
    def on_load():
        pm = load_scene_local(LOCAL_FILE_DEFAULT)
        pts_all, vids_all = stack_views(pm)
        pts_vx, vids_vx   = voxel_downsample_with_ids(pts_all, vids_all, VOXEL_SIZE)
        pts_vx, vids_vx   = hard_cap(pts_vx, vids_vx, DOWNSAMPLE_N_MAX)
        ps = adaptive_point_size(pts_vx.shape[0])
        fig3d = base_figure_gray_depth(pts_vx, ps, camera=DEFAULT_CAM)
        msg = f"✅ Loaded {os.path.basename(LOCAL_FILE_DEFAULT)} | Views: {pm.shape[0]} | Points: {pts_vx.shape[0]:,}"
        return fig3d, TOP_VIEW_IMAGE_PATH, msg, pm, pts_vx, vids_vx

    def on_submit(text, topk, pm, pts_vx, vids_vx):
        if pm is None:
            return gr.update(), TOP_VIEW_IMAGE_PATH, "⚠️ Scene not loaded yet."
        k = int(max(1, min(12, int(topk)))) if topk else TOPK_VIEWS_DEFAULT
        top_views = rank_views_for_text(model, text, pm, device, topk=k)
        ps = adaptive_point_size(pts_vx.shape[0])
        fig = highlight_views_3d(pts_vx, vids_vx, top_views, ps, camera=DEFAULT_CAM)
        msg = f"Highlighted views (top-{k}): {top_views}"
        return fig, TOP_VIEW_IMAGE_PATH, msg

    demo.load(on_load, inputs=[], outputs=[plot3d, img_ref, status, pm_state, pts_state, vids_state])
    submit_btn.click(on_submit, inputs=[text_in, topk_in, pm_state, pts_state, vids_state],
                     outputs=[plot3d, img_ref, status])

if __name__ == "__main__":
    demo.launch()