| """Per-pixel language targets (label_text v1: SigLIP2-base text embeddings of 624 concepts) for splatted features. |
| |
| Source: point4d-datasets tools/label_text/v1 (README there). Every semantic class / object category maps to one |
| concept row of `targets [624, 768]` (L2-normalised google/siglip2-base-patch16-224 text embeddings); row -1 means |
| IGNORE (unlabelled, void, background, ...), and ids or names missing from a table are IGNORE too. |
| * table key "semantic id": rows = lut[semantic] |
| * table key "category name": rows = classes[instances[str(instance)]["name"]]["row"] |
| * no semantic channel but a semantic-id table (data-012 in the HF mirror): read instances[...]["semantic_id"] |
| Loss (as recommended there): contrastive over all concepts, logits = cos(pred, targets) / tau with a learnable tau |
| initialised at the fitted 0.009; plain cosine regression separates classes poorly. |
| """ |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| DEFAULT_DIR = "/scratch/data/point4d-datasets/tools/label_text/v1" |
| IGNORE_ROW = -1 |
|
|
|
|
| class LabelText: |
| def __init__(self, root=DEFAULT_DIR): |
| self.root = Path(root) |
| z = np.load(self.root / "label_embeddings_v1.npz") |
| self.targets = z["targets"].astype(np.float32) |
| self.concepts = [str(c) for c in z["concept_names"]] |
| self.tables = {} |
| for path in self.root.glob("*_label_text.json"): |
| table = json.loads(path.read_text()) |
| classes = table.get("classes") or {} |
| if not classes: |
| continue |
| semantic = table.get("key", "").startswith("semantic") |
| entry = dict(semantic=semantic, rows={k: int(v.get("row", IGNORE_ROW)) for k, v in classes.items()}) |
| if semantic: |
| ids = [int(k) for k in classes] |
| lut = np.full(max(ids) + 1, IGNORE_ROW, np.int32) |
| for k, row in entry["rows"].items(): |
| lut[int(k)] = row |
| entry["lut"] = lut |
| self.tables[table.get("dataset", path.name.removesuffix("_label_text.json"))] = entry |
|
|
| def rows(self, dataset, semantic=None, instance=None, instances=None): |
| """(H, W) int32 concept rows for one frame, or None when the dataset has no usable channel/table.""" |
| table = self.tables.get(dataset) |
| if table is None: |
| return None |
| if table["semantic"]: |
| ids = semantic |
| if ids is None and instance is not None and instances: |
| lut = np.full(int(instance.max()) + 1, -1, np.int64) |
| for k, v in instances.items(): |
| if int(k) <= instance.max() and v.get("semantic_id") is not None: |
| lut[int(k)] = int(v["semantic_id"]) |
| ids = lut[instance] |
| if ids is None: |
| return None |
| ids = ids.astype(np.int64) |
| lut = table["lut"] |
| ok = (ids >= 0) & (ids < len(lut)) |
| return np.where(ok, lut[np.clip(ids, 0, len(lut) - 1)], IGNORE_ROW).astype(np.int32) |
| if instance is None or not instances: |
| return None |
| lut = np.full(int(instance.max()) + 1, IGNORE_ROW, np.int32) |
| for k, v in instances.items(): |
| name = v.get("name") |
| if name is not None and int(k) <= instance.max(): |
| lut[int(k)] = table["rows"].get(name, IGNORE_ROW) |
| return lut[instance] |
|
|
|
|
| class LanguageDecoder(nn.Module): |
| """Splatted per-pixel latent (lang_dim) -> 768-d SigLIP2 space; contrastive logits over the concept targets.""" |
|
|
| def __init__(self, lang_dim, targets, tau=0.009): |
| super().__init__() |
| self.proj = nn.Sequential(nn.Linear(lang_dim, 256), nn.GELU(), nn.Linear(256, targets.shape[1])) |
| self.register_buffer("targets", F.normalize(torch.as_tensor(targets, dtype=torch.float32), dim=-1), persistent=False) |
| self.log_tau = nn.Parameter(torch.tensor(float(np.log(tau)))) |
|
|
| def embed(self, latent): |
| return F.normalize(self.proj(latent.float()), dim=-1) |
|
|
| def loss(self, latent, rows, max_pixels=4096): |
| """Cross-entropy of cos(embed, targets) / tau against concept rows (IGNORE skipped); also top-1 accuracy.""" |
| valid = rows >= 0 |
| if not valid.any(): |
| return latent.sum() * 0.0, {} |
| lat, tgt = latent[valid], rows[valid].long() |
| if len(tgt) > max_pixels: |
| pick = torch.randperm(len(tgt), device=tgt.device)[:max_pixels] |
| lat, tgt = lat[pick], tgt[pick] |
| logits = self.embed(lat) @ self.targets.T / self.log_tau.exp() |
| loss = F.cross_entropy(logits, tgt) |
| return loss, dict(lang=float(loss.detach()), lang_acc=float((logits.argmax(-1) == tgt).float().mean()), |
| lang_tau=float(self.log_tau.exp().detach())) |
|
|
