| """Jev-1.14-style `box` questions: bounding box of a UI element as fractions of the image.
|
|
|
| Answer format (same as Jev 1.14):
|
| {"type": "box", "box": {"top": .., "bottom": .., "left": .., "right": ..}, "confidence": 0..1}
|
| top/bottom are fractions of the image height, left/right fractions of the width, origin top-left,
|
| left < right and top < bottom. There is no "not found" answer: presence is a separate noul question
|
| (convention: `<box qid>_visible`).
|
|
|
| Edge layout of every tensor in this module: (..., 4) = (top, bottom, left, right).
|
|
|
| Head: for each edge, logits over K bins (distribution focal loss style, GFL/YOLOv8).
|
| Bin i has centre i / (K - 1), so 0 and 1 are exactly representable. Decoded edge = expectation
|
| over bin centres; ordering is enforced at decode with min/max per axis.
|
| Confidence: a learned sigmoid head regressing the IoU with gold (BCE against the detached IoU).
|
| """
|
| from __future__ import annotations
|
| import math
|
| from dataclasses import dataclass
|
| from typing import Any
|
| import torch
|
| import torch.nn.functional as F
|
|
|
| BOX_EDGES = ("top", "bottom", "left", "right")
|
| DEFAULT_BINS = 64
|
| FOURIER_BANDS = 8
|
|
|
|
|
| def is_box_question(question: dict[str, Any]) -> bool:
|
| return str(question.get("type")) == "box"
|
|
|
|
|
|
|
|
|
| def bin_centers(bins: int, device=None, dtype=torch.float32) -> torch.Tensor:
|
| return torch.linspace(0.0, 1.0, bins, device=device, dtype=dtype)
|
|
|
|
|
| def two_hot(target: torch.Tensor, bins: int) -> torch.Tensor:
|
| """Soft label on the two bins adjacent to each target value in [0, 1]. (...,) -> (..., bins)."""
|
| t = target.float().clamp(0.0, 1.0) * (bins - 1)
|
| lo = t.floor().long().clamp(max=bins - 2)
|
| w_hi = t - lo.float()
|
| out = torch.zeros(*target.shape, bins, device=target.device)
|
| out.scatter_(-1, lo.unsqueeze(-1), (1.0 - w_hi).unsqueeze(-1))
|
| out.scatter_add_(-1, (lo + 1).unsqueeze(-1), w_hi.unsqueeze(-1))
|
| return out
|
|
|
|
|
| def decode_edges(logits: torch.Tensor) -> torch.Tensor:
|
| """Expectation over bin centres. (..., 4, K) logits -> (..., 4) edges (unordered)."""
|
| p = logits.float().softmax(-1)
|
| return (p * bin_centers(logits.shape[-1], logits.device)).sum(-1)
|
|
|
|
|
| def order_box(edges: torch.Tensor) -> torch.Tensor:
|
| """Enforce top <= bottom and left <= right by swapping (min/max per axis; differentiable)."""
|
| t, b, l, r = edges.unbind(-1)
|
| return torch.stack([torch.minimum(t, b), torch.maximum(t, b), torch.minimum(l, r), torch.maximum(l, r)], -1)
|
|
|
|
|
| def decode_box(logits: torch.Tensor) -> torch.Tensor:
|
| return order_box(decode_edges(logits))
|
|
|
|
|
| def max_prob_confidence(logits: torch.Tensor) -> torch.Tensor:
|
| """Fallback confidence: mean over the 4 edges of the sum of the two most likely adjacent bins."""
|
| p = logits.float().softmax(-1)
|
| pair = p[..., :-1] + p[..., 1:]
|
| return pair.max(-1).values.mean(-1)
|
|
|
|
|
|
|
|
|
| def box_iou(a: torch.Tensor, b: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
|
| """IoU of (top, bottom, left, right) boxes, broadcasting over leading dims."""
|
| ih = (torch.minimum(a[..., 1], b[..., 1]) - torch.maximum(a[..., 0], b[..., 0])).clamp(min=0)
|
| iw = (torch.minimum(a[..., 3], b[..., 3]) - torch.maximum(a[..., 2], b[..., 2])).clamp(min=0)
|
| inter = ih * iw
|
| area_a = (a[..., 1] - a[..., 0]).clamp(min=0) * (a[..., 3] - a[..., 2]).clamp(min=0)
|
| area_b = (b[..., 1] - b[..., 0]).clamp(min=0) * (b[..., 3] - b[..., 2]).clamp(min=0)
|
| return inter / (area_a + area_b - inter + eps)
|
|
|
|
|
| def generalized_iou(a: torch.Tensor, b: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
|
| """GIoU in [-1, 1] of (top, bottom, left, right) boxes (both must be ordered)."""
|
| ih = (torch.minimum(a[..., 1], b[..., 1]) - torch.maximum(a[..., 0], b[..., 0])).clamp(min=0)
|
| iw = (torch.minimum(a[..., 3], b[..., 3]) - torch.maximum(a[..., 2], b[..., 2])).clamp(min=0)
|
| inter = ih * iw
|
| area_a = (a[..., 1] - a[..., 0]) * (a[..., 3] - a[..., 2])
|
| area_b = (b[..., 1] - b[..., 0]) * (b[..., 3] - b[..., 2])
|
| union = area_a + area_b - inter
|
| hull = ((torch.maximum(a[..., 1], b[..., 1]) - torch.minimum(a[..., 0], b[..., 0]))
|
| * (torch.maximum(a[..., 3], b[..., 3]) - torch.minimum(a[..., 2], b[..., 2])))
|
| return inter / (union + eps) - (hull - union) / (hull + eps)
|
|
|
|
|
|
|
|
|
| def dfl_loss(logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
| """Distribution focal loss: soft-label CE on the two bins adjacent to each edge, mean over edges.
|
| logits (..., 4, K), target (..., 4) in [0, 1]."""
|
| logp = logits.float().log_softmax(-1)
|
| return -(two_hot(target, logits.shape[-1]) * logp).sum(-1).mean()
|
|
|
|
|
| @dataclass
|
| class BoxLossWeights:
|
| dfl: float = 1.0
|
| l1: float = 5.0
|
| giou: float = 2.0
|
| conf: float = 1.0
|
|
|
|
|
| def box_loss(logits: torch.Tensor, conf_logit: torch.Tensor, target: torch.Tensor,
|
| weights: BoxLossWeights | None = None) -> tuple[torch.Tensor, dict[str, float]]:
|
| """Loss for one or more box questions. logits (N, 4, K), conf_logit (N,), target (N, 4) ordered gold.
|
| Returns mean over N and the (float) parts."""
|
| w = weights or BoxLossWeights()
|
| target = target.float()
|
| pred = decode_box(logits)
|
| dfl = dfl_loss(logits, target)
|
| l1 = (pred - target).abs().mean()
|
| giou = (1.0 - generalized_iou(pred, target)).mean()
|
| iou = box_iou(pred, target).detach()
|
| conf = F.binary_cross_entropy_with_logits(conf_logit.float(), iou)
|
| loss = w.dfl * dfl + w.l1 * l1 + w.giou * giou + w.conf * conf
|
| return loss, {"dfl": float(dfl.detach()), "l1": float(l1.detach()), "giou": float(giou.detach()),
|
| "conf": float(conf.detach()), "iou": float(iou.mean())}
|
|
|
|
|
|
|
|
|
| def box_answer(box: torch.Tensor | list[float], confidence: float, min_size: float = 1e-3, digits: int = 4) -> dict[str, Any]:
|
| """Jev-1.14 box answer. Guarantees 0 <= top < bottom <= 1 and 0 <= left < right <= 1."""
|
| t, b, l, r = [float(v) for v in (box.tolist() if torch.is_tensor(box) else box)]
|
| t, b = sorted((min(max(t, 0.0), 1.0), min(max(b, 0.0), 1.0)))
|
| l, r = sorted((min(max(l, 0.0), 1.0), min(max(r, 0.0), 1.0)))
|
|
|
| def widen(lo, hi):
|
| if hi - lo >= min_size:
|
| return lo, hi
|
| c = min(max((lo + hi) / 2, min_size / 2), 1 - min_size / 2)
|
| return c - min_size / 2, c + min_size / 2
|
| t, b = widen(t, b); l, r = widen(l, r)
|
| rd = lambda v: round(v, digits)
|
| return {"type": "box", "box": {"top": rd(t), "bottom": rd(b), "left": rd(l), "right": rd(r)},
|
| "confidence": rd(min(max(float(confidence), 0.0), 1.0))}
|
|
|
|
|
| def ece(confs, hits, bins=15) -> float:
|
| tot, n = 0.0, len(confs)
|
| for k in range(bins):
|
| lo, hi = k / bins, (k + 1) / bins
|
| idx = [i for i, c in enumerate(confs) if lo < c <= hi or (k == 0 and c == 0)]
|
| if idx:
|
| tot += len(idx) / n * abs(sum(hits[i] for i in idx) / len(idx) - sum(confs[i] for i in idx) / len(idx))
|
| return tot
|
|
|
|
|
| def box_metrics(ious: list[float], confs: list[float], maxprob_confs: list[float] | None = None, bins: int = 15) -> dict[str, Any]:
|
| n = len(ious)
|
| if not n:
|
| return {"n_box_questions": 0}
|
| hit50 = [int(i >= 0.5) for i in ious]
|
| out = {"n_box_questions": n, "mean_iou": sum(ious) / n, "iou_ge_0.5": sum(hit50) / n,
|
| "iou_ge_0.75": sum(int(i >= 0.75) for i in ious) / n, "mean_confidence": sum(confs) / n,
|
| f"ece_{bins}bin_conf_vs_iou50": ece(confs, hit50, bins)}
|
| if maxprob_confs:
|
| out[f"ece_{bins}bin_maxprob_vs_iou50"] = ece(maxprob_confs, hit50, bins)
|
| return out
|
|
|
|
|
|
|
|
|
| def fourier_features(yx: torch.Tensor, bands: int = FOURIER_BANDS) -> torch.Tensor:
|
| """(N, 2) coords in [0, 1] -> (N, 4 * bands) sin/cos features."""
|
| freqs = (2.0 ** torch.arange(bands, device=yx.device, dtype=torch.float32)) * math.pi
|
| ang = yx.float().unsqueeze(-1) * freqs
|
| return torch.cat([ang.sin(), ang.cos()], -1).flatten(1)
|
|
|
|
|
| @dataclass
|
| class BoxPrediction:
|
| logits: torch.Tensor
|
| conf_logit: torch.Tensor
|
| question_ids: tuple[str, ...]
|
|
|
|
|
| class BoxHead(torch.nn.Module):
|
| """Box head on the backbone's last hidden states.
|
|
|
| Query per box question = projection of the same pooled question representation the joint schema
|
| head uses (mean of the normalised hidden states over the instruction span) + the global (last-token)
|
| vector, layer-normed and expanded to 4 edge queries (top/bottom/left/right) with a learned edge
|
| embedding of the same (unit) scale, so the edges are distinguishable from the first step.
|
| The edge queries run through pre-norm transformer decoder layers (self-attention between the 4 edges
|
| + cross-attention to the projected sequence). Image tokens of the memory get an explicit 2D Fourier
|
| position encoding (fractional row/col of the merged patch) so the head can read *where* evidence is.
|
| Each edge query -> K bin logits through its own (per-edge) scorer; the 4 edge queries + per-edge (max prob, entropy) -> IoU logit.
|
| """
|
|
|
| def __init__(self, hidden_size: int, width: int = 256, layers: int = 2, heads: int = 4,
|
| feedforward: int = 1024, bins: int = DEFAULT_BINS, fourier_bands: int = FOURIER_BANDS,
|
| dropout: float = 0.0) -> None:
|
| super().__init__()
|
| self.bins, self.fourier_bands = bins, fourier_bands
|
| self.hidden_norm = torch.nn.LayerNorm(hidden_size)
|
| self.memory_projection = torch.nn.Linear(hidden_size, width, bias=False)
|
| self.position_projection = torch.nn.Linear(4 * fourier_bands, width)
|
| self.question_projection = torch.nn.Linear(hidden_size, width, bias=False)
|
| self.global_projection = torch.nn.Linear(hidden_size, width, bias=False)
|
| self.query_norm = torch.nn.LayerNorm(width)
|
| self.edge_embedding = torch.nn.Parameter(torch.randn(4, width))
|
| self.layers = torch.nn.ModuleList([
|
| torch.nn.TransformerDecoderLayer(d_model=width, nhead=heads, dim_feedforward=feedforward, dropout=dropout,
|
| activation="gelu", batch_first=True, norm_first=True)
|
| for _ in range(layers)])
|
| self.output_norm = torch.nn.LayerNorm(width)
|
| bound = 1.0 / math.sqrt(width)
|
| self.bin_weight = torch.nn.Parameter(torch.empty(4, width, bins).uniform_(-bound, bound))
|
| self.bin_bias = torch.nn.Parameter(torch.empty(4, bins).uniform_(-bound, bound))
|
| self.confidence_scorer = torch.nn.Sequential(
|
| torch.nn.Linear(4 * width + 8, width), torch.nn.GELU(), torch.nn.Dropout(dropout), torch.nn.Linear(width, 1))
|
|
|
| def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, records: list[Any]) -> list[BoxPrediction | None]:
|
| results: list[BoxPrediction | None] = []
|
| normalized_hidden = self.hidden_norm(hidden_states)
|
| for batch_index, record in enumerate(records):
|
| box_questions = getattr(record, "box_questions", ()) or ()
|
| if not box_questions:
|
| results.append(None)
|
| continue
|
| sequence_length = int(attention_mask[batch_index].sum().item())
|
| sequence_hidden = normalized_hidden[batch_index, :sequence_length]
|
| memory = self.memory_projection(sequence_hidden)
|
| coords = getattr(record, "image_coords", None)
|
| if coords is not None and len(coords[0]):
|
| index = torch.tensor(coords[0], dtype=torch.long, device=memory.device)
|
| yx = torch.tensor(coords[1], dtype=torch.float32, device=memory.device)
|
| position = self.position_projection(fourier_features(yx, self.fourier_bands).to(memory.dtype))
|
| memory = memory.index_add(0, index, position)
|
| question_vectors = torch.stack([sequence_hidden[s:e].mean(0) for s, e in (q.question_span for q in box_questions)])
|
| base = self.query_norm(self.question_projection(question_vectors) + self.global_projection(sequence_hidden[-1]).unsqueeze(0))
|
| queries = (base.unsqueeze(1) + self.edge_embedding.to(base.dtype)).reshape(1, -1, base.shape[-1])
|
|
|
| q_index = torch.arange(len(box_questions), device=memory.device).repeat_interleave(4)
|
| self_mask = q_index.unsqueeze(0) != q_index.unsqueeze(1)
|
| for layer in self.layers:
|
| queries = layer(queries, memory.unsqueeze(0), tgt_mask=self_mask)
|
| queries = self.output_norm(queries[0]).view(len(box_questions), 4, -1)
|
| logits = torch.einsum("qew,ewk->qek", queries, self.bin_weight.to(queries.dtype)) + self.bin_bias.to(queries.dtype)
|
| p = logits.detach().float().softmax(-1)
|
| stats = torch.cat([p.max(-1).values, -(p * p.clamp_min(1e-9).log()).sum(-1) / math.log(self.bins)], -1)
|
| conf_logit = self.confidence_scorer(torch.cat([queries.flatten(1), stats.to(queries.dtype)], -1)).squeeze(-1)
|
| results.append(BoxPrediction(logits, conf_logit, tuple(q.question_id for q in box_questions)))
|
| return results
|
|
|
