| |
| """C08 outputs: the result classes of the 13 bundles and their ``POST /predict`` JSON (BUNDLE_CONVENTIONS.md 7.4). |
| |
| Every class has ``to_dicts()`` (one plain dict per detection / pose / class) and ``to_dict(output_format)`` (the |
| whole response body: ``model``, ``frame_id``, the payload, ``meta``, ``timing_ms``). The Python API returns these |
| objects and the server returns ``to_dict(...)`` of the same object, so both agree bit for bit. |
| |
| ================== ================================= ================================================== |
| class families payload |
| ================== ================================= ================================================== |
| ``Detections3D`` CenterPoint, TransFusion, BEVFusion, ``detections`` [label, label_id, score, center, size, |
| PointPainting, StreamPETR, BEVDet, yaw, velocity], ``num_detections``; rows sorted by |
| BEVFormer, PTv3-det descending score at construction |
| ``Detections2D`` YOLOX ``detections`` [label, label_id, score, box_xyxy] + |
| ``extras`` (e.g. ``{"semseg": Mask2D}``) |
| ``Segmentation3D`` FRNet, PTv3-seg ``labels`` (npz uint8/uint16 [N]), ``scores``, |
| ``class_names``, ``class_counts`` |
| ``Mask2D`` SceneSeg (YOLOX semseg extra) ``mask`` (png or npz), ``class_names``, ``class_counts`` |
| ``Trajectory`` Diffusion Planner ``trajectory`` [[x, y, yaw, ...] x T], ``columns``, |
| ``turn_indicator``, ``predicted_agents`` |
| ================== ================================= ================================================== |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
| from typing import Any, Dict, List, Optional, Sequence |
|
|
| import numpy as np |
|
|
| from .io import encode_array, encode_png, to_jsonable |
|
|
| __all__ = ["Detections3D", "Detections2D", "Segmentation3D", "Mask2D", "Trajectory", "label_name"] |
|
|
|
|
| def label_name(labels: Sequence[str], label_id: int) -> str: |
| """``labels[label_id]``, or the id as text when it is out of range.""" |
| return labels[label_id] if 0 <= label_id < len(labels) else str(label_id) |
|
|
|
|
| def _envelope(model: str, frame_id: str, meta: dict, timing_ms: dict) -> Dict[str, Any]: |
| return {"model": model, "frame_id": frame_id, "meta": to_jsonable(meta), |
| "timing_ms": {k: round(float(v), 3) for k, v in timing_ms.items()}} |
|
|
|
|
| def _sorted_by_score(scores: np.ndarray) -> np.ndarray: |
| return np.argsort(-scores, kind="stable") |
|
|
|
|
| @dataclass |
| class Detections3D: |
| """3D boxes in ``frame_id`` (Autoware ``base_link``: x forward, y left, z up; metres, radians). |
| |
| ``boxes`` float32 ``[N, 7]`` = x, y, z (box centre), length, width, height, yaw; ``scores`` ``[N]``; |
| ``label_ids`` int32 ``[N]`` (index into ``labels``); ``velocities`` ``[N, 2]`` (vx, vy) m/s or None. |
| Rows are sorted by descending score (stable) at construction.""" |
|
|
| boxes: np.ndarray |
| scores: np.ndarray |
| label_ids: np.ndarray |
| velocities: Optional[np.ndarray] = None |
| labels: Sequence[str] = () |
| model: str = "" |
| frame_id: str = "base_link" |
| timing_ms: dict = field(default_factory=dict) |
| meta: dict = field(default_factory=dict) |
|
|
| def __post_init__(self) -> None: |
| self.boxes = np.asarray(self.boxes, np.float32).reshape(-1, 7) |
| self.scores = np.asarray(self.scores, np.float32).reshape(-1) |
| self.label_ids = np.asarray(self.label_ids, np.int32).reshape(-1) |
| n = len(self.scores) |
| if len(self.boxes) != n or len(self.label_ids) != n: |
| raise ValueError(f"boxes {self.boxes.shape}, scores {self.scores.shape}, labels {self.label_ids.shape}") |
| if self.velocities is not None: |
| self.velocities = np.asarray(self.velocities, np.float32).reshape(n, 2) |
| order = _sorted_by_score(self.scores) |
| self.boxes, self.scores, self.label_ids = self.boxes[order], self.scores[order], self.label_ids[order] |
| if self.velocities is not None: |
| self.velocities = self.velocities[order] |
|
|
| def __len__(self) -> int: |
| return int(self.scores.shape[0]) |
|
|
| @property |
| def label_names(self) -> List[str]: |
| return [label_name(self.labels, i) for i in self.label_ids.tolist()] |
|
|
| def to_dicts(self) -> List[Dict[str, Any]]: |
| out = [] |
| for i, name in enumerate(self.label_names): |
| x, y, z, length, width, height, yaw = (float(v) for v in self.boxes[i]) |
| d = {"label": name, "label_id": int(self.label_ids[i]), "score": round(float(self.scores[i]), 4), |
| "center": [round(x, 3), round(y, 3), round(z, 3)], |
| "size": [round(length, 3), round(width, 3), round(height, 3)], "yaw": round(yaw, 4)} |
| if self.velocities is not None: |
| d["velocity"] = [round(float(v), 3) for v in self.velocities[i]] |
| out.append(d) |
| return out |
|
|
| def to_dict(self, output_format: str = "json") -> Dict[str, Any]: |
| body = _envelope(self.model, self.frame_id, self.meta, self.timing_ms) |
| body.update(num_detections=len(self), detections=self.to_dicts()) |
| if output_format == "npz": |
| body["arrays"] = {"boxes": encode_array(self.boxes, key="boxes"), |
| "scores": encode_array(self.scores, key="scores"), |
| "label_ids": encode_array(self.label_ids, key="label_ids")} |
| if self.velocities is not None: |
| body["arrays"]["velocities"] = encode_array(self.velocities, key="velocities") |
| return body |
|
|
|
|
| @dataclass |
| class Mask2D: |
| """A per-pixel class map (H, W) (uint8 for <= 256 classes) in the source image's pixel grid.""" |
|
|
| mask: np.ndarray |
| class_names: Sequence[str] = () |
| model: str = "" |
| frame_id: str = "camera" |
| encoding: str = "png" |
| timing_ms: dict = field(default_factory=dict) |
| meta: dict = field(default_factory=dict) |
|
|
| def __post_init__(self) -> None: |
| self.mask = np.asarray(self.mask) |
| if self.mask.ndim != 2: |
| raise ValueError(f"mask must be (H, W), got {self.mask.shape}") |
| if self.encoding not in ("png", "npz"): |
| raise ValueError("encoding must be 'png' or 'npz'") |
|
|
| def class_counts(self) -> Dict[str, int]: |
| ids, counts = np.unique(self.mask, return_counts=True) |
| return {label_name(self.class_names, int(i)): int(c) for i, c in zip(ids, counts)} |
|
|
| def to_dicts(self) -> List[Dict[str, Any]]: |
| ids, counts = np.unique(self.mask, return_counts=True) |
| return [{"label": label_name(self.class_names, int(i)), "label_id": int(i), "pixels": int(c)} |
| for i, c in zip(ids, counts)] |
|
|
| def payload(self, output_format: str = "json") -> Dict[str, Any]: |
| """The encoded mask alone (also used when the mask is an extra of another output).""" |
| if self.encoding == "png" and output_format != "npz" and self.mask.dtype == np.uint8: |
| return encode_png(self.mask, key="mask") |
| return encode_array(self.mask, key="mask") |
|
|
| def to_dict(self, output_format: str = "json") -> Dict[str, Any]: |
| body = _envelope(self.model, self.frame_id, self.meta, self.timing_ms) |
| body.update(mask=self.payload(output_format), class_names=list(self.class_names), |
| class_counts=self.class_counts()) |
| return body |
|
|
|
|
| @dataclass |
| class Detections2D: |
| """2-D boxes in original image pixels: ``boxes_xyxy`` float32 ``[N, 4]``, ``scores``, ``label_ids``. |
| ``extras`` holds companion outputs encoded into the body by name (``{"semseg": Mask2D(...)}``). |
| Rows are sorted by descending score (stable) at construction.""" |
|
|
| boxes_xyxy: np.ndarray |
| scores: np.ndarray |
| label_ids: np.ndarray |
| labels: Sequence[str] = () |
| extras: Dict[str, Any] = field(default_factory=dict) |
| model: str = "" |
| frame_id: str = "camera" |
| timing_ms: dict = field(default_factory=dict) |
| meta: dict = field(default_factory=dict) |
|
|
| def __post_init__(self) -> None: |
| self.boxes_xyxy = np.asarray(self.boxes_xyxy, np.float32).reshape(-1, 4) |
| self.scores = np.asarray(self.scores, np.float32).reshape(-1) |
| self.label_ids = np.asarray(self.label_ids, np.int32).reshape(-1) |
| if not (len(self.boxes_xyxy) == len(self.scores) == len(self.label_ids)): |
| raise ValueError("boxes_xyxy, scores and label_ids lengths differ") |
| order = _sorted_by_score(self.scores) |
| self.boxes_xyxy, self.scores, self.label_ids = self.boxes_xyxy[order], self.scores[order], self.label_ids[order] |
|
|
| def __len__(self) -> int: |
| return int(self.scores.shape[0]) |
|
|
| def to_dicts(self) -> List[Dict[str, Any]]: |
| return [{"label": label_name(self.labels, int(self.label_ids[i])), "label_id": int(self.label_ids[i]), |
| "score": round(float(self.scores[i]), 4), |
| "box_xyxy": [round(float(v), 2) for v in self.boxes_xyxy[i]]} for i in range(len(self))] |
|
|
| def to_dict(self, output_format: str = "json") -> Dict[str, Any]: |
| body = _envelope(self.model, self.frame_id, self.meta, self.timing_ms) |
| body.update(num_detections=len(self), detections=self.to_dicts()) |
| for name, extra in self.extras.items(): |
| if hasattr(extra, "payload"): |
| body[name] = extra.payload(output_format) |
| elif isinstance(extra, np.ndarray): |
| body[name] = encode_array(extra, key=name) |
| else: |
| body[name] = to_jsonable(extra) |
| if output_format == "npz": |
| body["arrays"] = {"boxes_xyxy": encode_array(self.boxes_xyxy, key="boxes_xyxy"), |
| "scores": encode_array(self.scores, key="scores"), |
| "label_ids": encode_array(self.label_ids, key="label_ids")} |
| return body |
|
|
|
|
| @dataclass |
| class Segmentation3D: |
| """Per-point classes, in input point order after NaN removal (stated in SERVING.md). ``label_ids`` ``[N]``; |
| ``scores`` optional ``[N]`` (winning-class probability) or ``[N, C]``.""" |
|
|
| label_ids: np.ndarray |
| class_names: Sequence[str] = () |
| scores: Optional[np.ndarray] = None |
| model: str = "" |
| frame_id: str = "base_link" |
| timing_ms: dict = field(default_factory=dict) |
| meta: dict = field(default_factory=dict) |
|
|
| def __post_init__(self) -> None: |
| ids = np.asarray(self.label_ids).reshape(-1) |
| if ids.size and ids.min() < 0: |
| raise ValueError("label ids must be >= 0") |
| self.label_ids = ids.astype(np.uint8 if (not ids.size or ids.max() < 256) else np.uint16) |
| if self.scores is not None: |
| self.scores = np.asarray(self.scores, np.float32) |
| if self.scores.shape[0] != ids.shape[0]: |
| raise ValueError("scores and label_ids lengths differ") |
|
|
| def __len__(self) -> int: |
| return int(self.label_ids.shape[0]) |
|
|
| def class_counts(self) -> Dict[str, int]: |
| counts = np.bincount(self.label_ids.astype(np.int64), minlength=len(self.class_names)) |
| return {label_name(self.class_names, i): int(c) for i, c in enumerate(counts) if c or i < len(self.class_names)} |
|
|
| def to_dicts(self) -> List[Dict[str, Any]]: |
| counts = np.bincount(self.label_ids.astype(np.int64), minlength=len(self.class_names)) |
| return [{"label": label_name(self.class_names, i), "label_id": i, "points": int(c)} |
| for i, c in enumerate(counts)] |
|
|
| def to_dict(self, output_format: str = "json") -> Dict[str, Any]: |
| body = _envelope(self.model, self.frame_id, self.meta, self.timing_ms) |
| body.update(num_points=len(self), labels=encode_array(self.label_ids, key="labels"), |
| class_names=list(self.class_names), class_counts=self.class_counts()) |
| if self.scores is not None: |
| body["scores"] = encode_array(self.scores, key="scores") |
| return body |
|
|
|
|
| @dataclass |
| class Trajectory: |
| """A planned trajectory ``poses`` ``[T, D]`` whose columns are named by ``columns`` (default x, y, yaw) in |
| ``frame_id``; optional ``turn_indicator`` and ``predicted_agents`` ``[A, T, D']``.""" |
|
|
| poses: np.ndarray |
| columns: Sequence[str] = ("x", "y", "yaw") |
| turn_indicator: Any = None |
| predicted_agents: Optional[np.ndarray] = None |
| model: str = "" |
| frame_id: str = "base_link" |
| timing_ms: dict = field(default_factory=dict) |
| meta: dict = field(default_factory=dict) |
|
|
| def __post_init__(self) -> None: |
| self.poses = np.asarray(self.poses, np.float32) |
| if self.poses.ndim != 2 or self.poses.shape[1] != len(self.columns): |
| raise ValueError(f"poses {self.poses.shape} do not match columns {list(self.columns)}") |
| if self.predicted_agents is not None: |
| self.predicted_agents = np.asarray(self.predicted_agents, np.float32) |
|
|
| def __len__(self) -> int: |
| return int(self.poses.shape[0]) |
|
|
| def to_dicts(self) -> List[Dict[str, Any]]: |
| return [{c: round(float(v), 4) for c, v in zip(self.columns, row)} for row in self.poses] |
|
|
| def to_dict(self, output_format: str = "json") -> Dict[str, Any]: |
| body = _envelope(self.model, self.frame_id, self.meta, self.timing_ms) |
| body.update(num_poses=len(self), columns=list(self.columns), |
| trajectory=[[round(float(v), 4) for v in row] for row in self.poses], |
| turn_indicator=to_jsonable(self.turn_indicator)) |
| if self.predicted_agents is not None: |
| body["predicted_agents"] = encode_array(self.predicted_agents, key="predicted_agents") |
| if output_format == "npz": |
| body["arrays"] = {"poses": encode_array(self.poses, key="poses")} |
| return body |
|
|
