File size: 5,056 Bytes
9dffa63 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """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) # (C, 768), L2-normalised
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 # e.g. ADT: empty for now -> all IGNORE
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: # semantic id via instances.json
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()))
|