AntonioJun/workspace

1"""Model-output adapters -- the ONLY file that knows any model's native output shape.2 3Every adapter returns THE canonical geometry dict below. geometric.py sees only this dict:4 5    instances: class -> list of already-reconstructed 3D instance dicts6    stats:     class -> raw/merged/peak counts7    scene_pts: full-scene world points used by gravity/floor-area math8    cameras:   world-space camera positions (only to orient gravity upward)9    raw_inputs: optional raw depth/pose/mask bundle for exact reference-math execution10 11Adapters may decode depth, masks, queries, meshes, voxels, Gaussians, or anything else. None12of those representations cross this file boundary.13"""14 15from __future__ import annotations16 17import gzip18import os19import pickle20import sys21 22import numpy as np23 24CACHE_ROOT = os.environ.get("VSI_CACHE_ROOT", "/root/data/caches")25FPS = float(os.environ.get("VSI_FPS", "6"))26MODELS_ROOT = os.environ.get("VSI_MODELS_ROOT", "/root/models")27DA3_ROOT = os.environ.get("VSI_DA3_ROOT", os.path.join(MODELS_ROOT, "depth-anything-3"))28SEGVGGT_ROOT = os.environ.get("VSI_SEGVGGT_ROOT", os.path.join(MODELS_ROOT, "SegVGGT"))29 30 31RAW_CACHE_ROOT = CACHE_ROOT32SPATIAL_CODE_FORMATS = ("compact", "explicit")33 34 35def _raw_cache_root(root=None):36    """Resolve the shared inference-output root without importing encoder config."""37    return root or os.environ.get("VSI_CACHE_ROOT", RAW_CACHE_ROOT)38 39 40def validate_spatial_code_format(spatial_code_format):41    """Validate one output format at the model-agnostic geometry boundary."""42    if spatial_code_format not in SPATIAL_CODE_FORMATS:43        raise ValueError(44            f"unknown spatial-code format {spatial_code_format!r}; "45            f"expected {SPATIAL_CODE_FORMATS}"46        )47    return spatial_code_format48 49 50# ==========================================================================================51# CANONICAL SHAPE -- validation only. Plain dicts on purpose: no schema/contract class layer.52# ==========================================================================================53 54 55def _pts(x, where):56    p = np.asarray(x, np.float32)57    if p.ndim != 2 or p.shape[1] != 3:58        raise ValueError(f"{where} must be [N,3], got {p.shape}")59    return p[np.isfinite(p).all(1)]60 61 62class EmptySceneError(ValueError):63    """Raised when a valid adapter result contains no spatial instances."""64 65 66def validate(scene):67    """Coerces one adapter result to THE canonical shape and fails loudly on drift."""68    if not isinstance(scene, dict):69        raise TypeError("adapter result must be a dict")70    if not scene.get("instances"):71        raise EmptySceneError("adapter produced no instances")72    for cls, insts in scene["instances"].items():73        for i in insts:74            i["pts"] = _pts(i["pts"], f"{cls}.pts")75            i["best_pts"] = _pts(i.get("best_pts", i["pts"]), f"{cls}.best_pts")76            i["n"] = int(i.get("n", len(i["pts"])))77            i["frames"] = set(i.get("frames", ()))78            i["nframes"] = int(i.get("nframes", len(i["frames"])))79            i["first_time"] = float(i.get("first_time", 0.0))80            if i.get("conf") is not None:81                i["conf"] = np.asarray(i["conf"], np.float32).reshape(-1)82    scene["scene_pts"] = _pts(scene["scene_pts"], "scene_pts")83    if scene.get("cameras") is not None:84        scene["cameras"] = _pts(scene["cameras"], "cameras")85    return scene86 87 88# ==========================================================================================89# DA3 + SAM3 -- compatibility adapter for the existing cache. All depth/mask projection lives90# HERE, not geometric.py. The geometric cleanup/measurement math is unchanged downstream.91# ==========================================================================================92 93 94def _load_masks(path):95    candidates = [path]96    if os.path.isdir(path):97        candidates = [98            os.path.join(path, n)99            for n in ("sam3.pkl.gz", "sam3_full.pkl.gz", "sam3.pkl")100        ]101    for path in candidates:102        if not os.path.exists(path):103            continue104        op = gzip.open if path.endswith(".gz") else open105        with op(path, "rb") as f:106            packed = pickle.load(f)107        out = {}108        for cls, frames in packed.items():109            out[cls] = {}110            for fi, objects in frames.items():111                out[cls][int(fi)] = {}112                for oid, value in objects.items():113                    if isinstance(value, tuple) and len(value) == 2:114                        bits, shape = value115                        m = (116                            np.unpackbits(bits)[: int(np.prod(shape))]117                            .reshape(shape)118                            .astype(bool)119                        )120                    else:121                        m = np.asarray(value, bool)122                    out[cls][int(fi)][int(oid)] = m123        return out124    raise FileNotFoundError(f"no SAM3 cache found at {path}")125 126 127def _sam3_response_masks(response, frame_index, allow_empty):128    """Normalize one tracked or independent SAM3 response to object-id masks."""129    if not isinstance(response, dict):130        raise ValueError(131            f"invalid SAM3 response at frame {frame_index}; expected a dict"132        )133    outputs = (134        response.get("outputs")135        if isinstance(response.get("outputs"), dict)136        else response137    )138    masks = outputs.get("masks")139    if masks is None:140        masks = outputs.get("out_binary_masks")141    if masks is None:142        if allow_empty and not outputs:143            return {}144        raise ValueError(f"SAM3 response at frame {frame_index} has no masks")145    masks = (146        masks.detach().cpu().numpy() if hasattr(masks, "detach") else np.asarray(masks)147    )148    if masks.ndim == 2:149        masks = masks[None]150    if masks.ndim == 4 and masks.shape[1] == 1:151        masks = masks[:, 0]152    if masks.ndim != 3:153        raise ValueError(154            f"SAM3 masks at frame {frame_index} must be [N,H,W], got {masks.shape}"155        )156    object_ids = outputs.get("out_obj_ids")157    if object_ids is None:158        object_ids = range(len(masks))159    elif hasattr(object_ids, "detach"):160        object_ids = object_ids.detach().cpu().numpy()161    return {162        int(object_id): mask.astype(bool) for object_id, mask in zip(object_ids, masks)163    }164 165 166def _load_native_sam3(path):167    """Decode tracked or independent native SAM3 frame responses."""168    try:169        import torch170    except ImportError as exc:171        raise RuntimeError("PyTorch is required to read a raw SAM3 .pt cache") from exc172    responses = torch.load(path, map_location="cpu", weights_only=False)173    if isinstance(responses, dict):174        decoded = {}175        for class_name, class_responses in responses.items():176            if isinstance(class_responses, dict) and "stream" in class_responses:177                # Lossless tracking caches retain complete request responses and178                # streamed envelopes. Only the stream is needed for geometry.179                stream = class_responses["stream"]180                if not isinstance(stream, list):181                    raise ValueError(182                        f"invalid SAM3 stream for {class_name!r}; expected a list"183                    )184                decoded[str(class_name)] = {185                    int(response["frame_index"]): _sam3_response_masks(186                        response, int(response["frame_index"]), True187                    )188                    for response in stream189                    if isinstance(response, dict)190                    and response.get("frame_index") is not None191                }192                continue193            if not isinstance(class_responses, list):194                raise ValueError(195                    f"invalid SAM3 responses for {class_name!r}; expected a list"196                )197            decoded[str(class_name)] = {198                frame_index: _sam3_response_masks(response, frame_index, True)199                for frame_index, response in enumerate(class_responses)200            }201        return decoded202    if not isinstance(responses, list):203        raise ValueError(f"invalid SAM3 raw cache {path}; expected a list or dict")204    class_name = os.environ.get("VSI_SAM3_PROMPT", "object")205    return {206        class_name: {207            frame_index: _sam3_response_masks(response, frame_index, False)208            for frame_index, response in enumerate(responses)209        }210    }211 212 213def _load_native_da3(path):214    """Decode DA3's exact pickled Prediction object into projection inputs."""215    model_root = str(DA3_ROOT)216    source_root = os.path.join(model_root, "src")217    if source_root not in sys.path:218        sys.path.insert(0, source_root)219    with open(path, "rb") as stream:220        prediction = pickle.load(stream)221 222    def field(name, required=True):223        value = getattr(prediction, name, None)224        if value is None and isinstance(prediction, dict):225            value = prediction.get(name)226        if required and value is None:227            raise ValueError(f"DA3 Prediction in {path} has no {name}")228        return value229 230    depth = np.asarray(field("depth"), np.float32)231    intr = np.asarray(field("intrinsics"), np.float32)232    extr = np.asarray(field("extrinsics"), np.float32)233    if extr.shape[-2:] == (3, 4):234        homogeneous = np.broadcast_to(235            np.eye(4, dtype=np.float32), extr.shape[:-2] + (4, 4)236        ).copy()237        homogeneous[..., :3, :] = extr238        extr = homogeneous239    if extr.shape[-2:] != (4, 4):240        raise ValueError(f"DA3 extrinsics must end in [3,4] or [4,4], got {extr.shape}")241    c2w = np.linalg.inv(extr).astype(np.float32)242    conf_value = field("conf", required=False)243    conf = np.asarray(conf_value, np.float32) if conf_value is not None else None244    return depth, intr, c2w, conf245 246 247def _resize_mask_nearest(mask, shape):248    """Resize one boolean mask to a depth-map shape with nearest sampling."""249    mask = np.asarray(mask, bool)250    if mask.ndim != 2:251        raise ValueError(f"mask must be two-dimensional, got {mask.shape}")252    if mask.shape == tuple(shape):253        return mask254    source_height, source_width = mask.shape255    target_height, target_width = shape256    rows = np.minimum(257        np.arange(target_height) * source_height // target_height,258        source_height - 1,259    )260    columns = np.minimum(261        np.arange(target_width) * source_width // target_width,262        source_width - 1,263    )264    return mask[np.ix_(rows, columns)]265 266 267def _backproject(depth, intrinsics, c2w, mask, conf=None):268    mask = _resize_mask_nearest(mask, depth.shape)269    ys, xs = np.nonzero(mask)270    z = depth[ys, xs]271    ok = np.isfinite(z) & (z > 0)272    ys, xs, z = ys[ok], xs[ok], z[ok]273    if not len(z):274        return np.zeros((0, 3), np.float32), None275    camera_points = np.stack(276        [277            (xs - intrinsics[0, 2]) * z / intrinsics[0, 0],278            (ys - intrinsics[1, 2]) * z / intrinsics[1, 1],279            z,280        ],281        1,282    )283    world_points = (c2w[:3, :3] @ camera_points.T).T + c2w[:3, 3]284    cf = conf[ys, xs].astype(np.float32) if conf is not None else None285    return world_points.astype(np.float32), cf286 287 288def _fusion_cache_paths(root, da3_path, sam3_path, scene):289    if scene:290        root = _raw_cache_root(root)291        da3_path = da3_path or os.path.join(root, "depth-anything-3", f"{scene}.pkl")292        sam3_path = sam3_path or os.path.join(root, "sam3", f"{scene}.pt")293    if da3_path is None:294        da3_path = os.path.join(root, "da3.npz") if root else None295    if sam3_path is None:296        sam3_path = root297    if not da3_path:298        raise ValueError("da3_path is required")299    return da3_path, sam3_path300 301 302def _load_fusion_inputs(da3_path, sam3_path):303    if str(da3_path).endswith(".pkl"):304        depth, intr, c2w, conf = _load_native_da3(da3_path)305        ft = np.arange(len(depth), dtype=np.float32) / FPS306    else:307        d = np.load(da3_path)308        depth, intr, c2w = d["depth"], d["intr"], d["c2w"]309        conf = d["conf"] if "conf" in d and d["conf"].size else None310        ft = (311            d["frame_times"]312            if "frame_times" in d313            else np.arange(len(depth), dtype=np.float32)314        )315    per = (316        _load_native_sam3(sam3_path)317        if str(sam3_path).endswith(".pt")318        else _load_masks(sam3_path)319    )320    return depth, intr, c2w, conf, ft, per321 322 323def _fuse_mask_instances(per, depth, intr, c2w, conf, frame_times):324    instances, stats = {}, {}325    for cls, frames in per.items():326        by_id = {}327        peak = max((len(v) for v in frames.values()), default=0)328        for fi, objects in frames.items():329            for oid, mask in objects.items():330                pts, cf = _backproject(331                    depth[fi],332                    intr[fi],333                    c2w[fi],334                    mask,335                    conf[fi] if conf is not None else None,336                )337                if not len(pts):338                    continue339                r = by_id.setdefault(340                    oid,341                    {342                        "chunks": [],343                        "conf": [],344                        "frames": set(),345                        "first_time": float(frame_times[fi]),346                    },347                )348                r["chunks"].append(pts)349                r["frames"].add(fi)350                if cf is not None:351                    r["conf"].append(cf)352        raw = []353        for r in by_id.values():354            pts = np.concatenate(r["chunks"], 0)355            raw.append(356                {357                    "pts": pts,358                    "best_pts": max(r["chunks"], key=len),359                    "observations": list(r["chunks"]),360                    "conf": np.concatenate(r["conf"], 0) if r["conf"] else None,361                    "frames": r["frames"],362                    "nframes": len(r["frames"]),363                    "n": len(pts),364                    "first_time": r["first_time"],365                }366            )367        if raw:368            instances[cls] = raw369            stats[cls] = {"raw": len(raw), "merged": len(raw), "peak": peak}370    return instances, stats371 372 373def _sample_scene_points(depth, intr, c2w):374    scene_points = []375    for fi in range(0, len(depth), 3):376        m = np.zeros_like(depth[fi], bool)377        m[::8, ::8] = True378        pts, _ = _backproject(depth[fi], intr[fi], c2w[fi], m)379        if len(pts):380            scene_points.append(pts)381    return np.concatenate(scene_points, 0)382 383 384def adapt_sam3_depth_anything_3(385    root=None, da3_path=None, sam3_path=None, scene=None, **_386):387    """Fuse raw DA3 geometry and raw per-frame SAM3 masks into canonical geometry."""388    da3_path, sam3_path = _fusion_cache_paths(root, da3_path, sam3_path, scene)389    depth, intr, c2w, conf, frame_times, per = _load_fusion_inputs(da3_path, sam3_path)390    instances, stats = _fuse_mask_instances(per, depth, intr, c2w, conf, frame_times)391    return validate(392        {393            "instances": instances,394            "stats": stats,395            "scene_pts": _sample_scene_points(depth, intr, c2w),396            "cameras": c2w[:, :3, 3],397            "raw_inputs": {398                "depth": depth,399                "intr": intr,400                "c2w": c2w,401                "conf": conf,402                "ftimes": frame_times,403                "per": per,404            },405        }406    )407 408 409# ==========================================================================================410# SegVGGT -- reads one deliberately boring NPZ export and groups already-world-space411# points by instance mask. Model inference is intentionally outside the encoder package.412# ==========================================================================================413 414 415SEGVGGT_CLASSES = """wall|floor|chair|table|door|couch|cabinet|shelf|desk|office chair|bed|pillow|sink|picture|window|toilet|bookshelf|monitor|curtain|book|armchair|coffee table|box|refrigerator|lamp|kitchen cabinet|towel|clothes|tv|nightstand|counter|dresser|stool|cushion|plant|ceiling|bathtub|end table|dining table|keyboard|bag|backpack|toilet paper|printer|tv stand|whiteboard|blanket|shower curtain|trash can|closet|stairs|microwave|stove|shoe|computer tower|bottle|bin|ottoman|bench|board|washing machine|mirror|copier|basket|sofa chair|file cabinet|fan|laptop|shower|paper|person|paper towel dispenser|oven|blinds|rack|plate|blackboard|piano|suitcase|rail|radiator|recycling bin|container|wardrobe|soap dispenser|telephone""".split(416    "|"417)418 419 420def _decode_segvggt_raw(path):421    """Decode the official SegVGGT.forward tensor dictionary for this adapter."""422    import sys423 424    try:425        import torch426        import torch.nn.functional as functional427    except ImportError as exc:428        raise RuntimeError("PyTorch is required to read a SegVGGT .pt cache") from exc429 430    model_root = str(SEGVGGT_ROOT)431    if model_root not in sys.path:432        sys.path.insert(0, model_root)433    try:434        from eval.instance_eval_common import predict_by_feat_instance435        from segvggt.utils.pose_enc import pose_encoding_to_extri_intri436    except ImportError as exc:437        raise RuntimeError(438            "SegVGGT is required to decode its native prediction dictionary"439        ) from exc440 441    raw = torch.load(path, map_location="cpu", weights_only=False)442    required = {"world_points", "instance_maps", "instance_labels", "pose_enc"}443    if not isinstance(raw, dict) or not required.issubset(raw):444        missing = (445            sorted(required - set(raw)) if isinstance(raw, dict) else sorted(required)446        )447        raise ValueError(f"invalid SegVGGT raw cache {path}; missing keys: {missing}")448 449    logits = raw["instance_maps"][0]450    query_count, frame_total, height, width = logits.shape451    masks, label_ids, _ = predict_by_feat_instance(452        raw["instance_labels"][0],453        logits.reshape(query_count, -1),454        mask_thr=float(os.environ.get("VSI_MASK_THR", "0.4")),455        npoint_thr=1,456    )457    masks = masks.reshape(-1, frame_total, height, width).cpu().numpy()458    label_ids = label_ids.cpu().numpy()459    keep = [index for index, label in enumerate(label_ids) if int(label) >= 2]460 461    world = raw["world_points"][0].float()462    if tuple(world.shape[1:3]) != (height, width):463        world = functional.interpolate(464            world.permute(0, 3, 1, 2),465            (height, width),466            mode="nearest",467        ).permute(0, 2, 3, 1)468    world = world.cpu().numpy()469 470    if "images" in raw:471        image_size = raw["images"].shape[-2:]472    elif "depth" in raw:473        image_size = raw["depth"].shape[2:4]474    else:475        image_size = (height, width)476    extrinsics, _ = pose_encoding_to_extri_intri(raw["pose_enc"].float(), image_size)477    extrinsics = extrinsics[0].cpu().numpy()478    rotations = extrinsics[:, :3, :3]479    translations = extrinsics[:, :3, 3]480    cameras = -np.einsum("sji,sj->si", rotations, translations)481    labels = np.asarray(482        [483            (484                SEGVGGT_CLASSES[int(label_ids[index])]485                if int(label_ids[index]) < len(SEGVGGT_CLASSES)486                else f"class {int(label_ids[index])}"487            )488            for index in keep489        ],490        dtype=object,491    )492    return {493        "world_points": world,494        "instance_masks": masks[keep].astype(bool),495        "labels": labels,496        "camera_positions": cameras.astype(np.float32),497    }498 499 500def adapt_segvggt(root=None, path=None, scene=None, **_):501    """Translate a raw-preserving SegVGGT cache to canonical geometry."""502    if path is None and scene:503        root = _raw_cache_root(root)504        path = os.path.join(root, "segvggt", f"{scene}.pt")505    else:506        path = path or root507    if path is None:508        raise ValueError("SegVGGT cache path is required")509    if os.path.isdir(path):510        path = os.path.join(path, f"{scene}.pt" if scene else "geometry.pt")511    if not os.path.exists(path):512        raise FileNotFoundError(f"SegVGGT raw cache does not exist: {path}")513    if str(path).endswith(".pt"):514        d = _decode_segvggt_raw(path)515    else:516        d = np.load(path, allow_pickle=True)517    world, masks = (518        np.asarray(d["world_points"], np.float32),519        np.asarray(d["instance_masks"], bool),520    )521    labels = [str(x) for x in d["labels"].tolist()]522    ft = (523        d["frame_times"]524        if "frame_times" in d525        else np.arange(world.shape[0], dtype=np.float32)526    )527    instances = {}528    for cls, mask in zip(labels, masks):529        fs = np.flatnonzero(mask.reshape(mask.shape[0], -1).any(1))530        chunks = [world[f][mask[f]] for f in fs if mask[f].any()]531        chunks = [p[np.isfinite(p).all(1)] for p in chunks if len(p)]532        if not chunks:533            continue534        pts = np.concatenate(chunks, 0)535        instances.setdefault(cls, []).append(536            {537                "pts": pts,538                "best_pts": max(chunks, key=len),539                "observations": list(chunks),540                "conf": None,541                "frames": set(map(int, fs)),542                "nframes": len(fs),543                "n": len(pts),544                "first_time": float(ft[fs[0]]) if len(fs) else 0.0,545            }546        )547    stats = {}548    for cls, insts in instances.items():549        n = {}550        for i in insts:551            for f in i["frames"]:552                n[f] = n.get(f, 0) + 1553        stats[cls] = {554            "raw": len(insts),555            "merged": len(insts),556            "peak": max(n.values(), default=len(insts)),557        }558    return validate(559        {560            "instances": instances,561            "stats": stats,562            "scene_pts": world.reshape(-1, 3),563            "cameras": d["camera_positions"] if "camera_positions" in d else None,564        }565    )566 567 568RAW_ADAPTERS = {569    "sam3+depth-anything-3": adapt_sam3_depth_anything_3,570    "segvggt": adapt_segvggt,571}572 573 574def available_models():575    """Return raw model formats supported by the encoder."""576    return tuple(sorted(RAW_ADAPTERS))577 578 579def adapt(model, **raw_cache):580    """Dispatch one model's native cache to its isolated format adapter."""581    adapter = RAW_ADAPTERS.get(model)582    if adapter is None:583        raise KeyError(584            f"no raw encoder adapter for {model!r}; expected one of {available_models()}"585        )586    return validate(adapter(**raw_cache))587