File size: 1,760 Bytes
e6f80f3 | 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 | """Support pointer decoding: the ``support_head``'s logits (floor | wall | request slots) -> ``support_parent`` ids.
A request that declares its support (``support_parent``, or a hard ``on`` constraint merged in by
``validation.effective_support_requests``) keeps it. The others take their most probable parent that keeps the
support graph acyclic, the most confident slots first; floor and wall never close a cycle, so every slot gets one.
"""
from __future__ import annotations
import torch
from fastfill.v2.batch import SUPPORT_CLASSES
def _closes_cycle(parents: dict, child: str, parent: str) -> bool:
seen = set()
while parent is not None and parent not in seen:
if parent == child:
return True
seen.add(parent)
parent = parents.get(parent)
return False
def decode_support_parents(logits: torch.Tensor, requests: list[dict]) -> list[str]:
"""``logits`` (n, 2 + n') of one scene's ``n`` requests (``requests``: effective support requests, in slot
order) -> one parent per request: "floor", "wall" or another request's id."""
ids = [r["id"] for r in requests]
names = list(SUPPORT_CLASSES) + ids
probability = logits[:, :len(names)].float().softmax(-1)
parents = {r["id"]: r["support_parent"] for r in requests if r.get("support_parent") is not None}
order = sorted((i for i, r in enumerate(requests) if r.get("support_parent") is None),
key=lambda i: -float(probability[i].max()))
for i in order:
for k in probability[i].argsort(descending=True).tolist():
if names[k] != ids[i] and not _closes_cycle(parents, ids[i], names[k]):
parents[ids[i]] = names[k]
break
return [parents[i] for i in ids]
|