128 lines
4.2 KiB
Python
128 lines
4.2 KiB
Python
"""Helpers for extracting GPR objects and locator observations from results."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
from python_app.models.dataset_model import ResultCollection, ResultPayload
|
|
|
|
|
|
def collection_payload_by_name(
|
|
collection: ResultCollection,
|
|
name: str,
|
|
kind: int | None = None,
|
|
) -> ResultPayload | None:
|
|
"""Return the first collection payload matching name and optional kind."""
|
|
for payload in collection.collection_payloads:
|
|
if payload.processing_name != name:
|
|
continue
|
|
if kind is not None and int(payload.kind) != int(kind):
|
|
continue
|
|
return payload
|
|
return None
|
|
|
|
|
|
def collection_payloads_by_prefix(
|
|
collection: ResultCollection,
|
|
prefix: str,
|
|
kind: int | None = None,
|
|
) -> list[ResultPayload]:
|
|
"""Return collection payloads matching a processing-name prefix."""
|
|
payloads: list[ResultPayload] = []
|
|
for payload in collection.collection_payloads:
|
|
if not str(payload.processing_name).startswith(prefix):
|
|
continue
|
|
if kind is not None and int(payload.kind) != int(kind):
|
|
continue
|
|
payloads.append(payload)
|
|
return payloads
|
|
|
|
|
|
def collection_has_gpr_payloads(collection: ResultCollection) -> bool:
|
|
"""Return whether collection carries GPR-specific collection payloads."""
|
|
return any(
|
|
str(payload.processing_name).startswith("gpr_")
|
|
for payload in collection.collection_payloads
|
|
)
|
|
|
|
|
|
def gpr_object_rows(collection: ResultCollection) -> np.ndarray:
|
|
"""Return object rows as `[x_m, z_m, score]` from a GPR collection."""
|
|
points_payload = collection_payload_by_name(collection, "gpr_points", kind=4)
|
|
if points_payload is not None:
|
|
points = np.asarray(points_payload.table, dtype=np.float32)
|
|
if points.ndim == 2 and points.shape[1] >= 3:
|
|
return points[:, :3]
|
|
|
|
centers_payload = collection_payload_by_name(collection, "gpr_region_centers", kind=4)
|
|
if centers_payload is not None:
|
|
centers = np.asarray(centers_payload.table, dtype=np.float32)
|
|
if centers.ndim == 2 and centers.shape[1] >= 3:
|
|
return centers[:, :3]
|
|
|
|
return np.zeros((0, 3), dtype=np.float32)
|
|
|
|
|
|
def locator_observations_from_collection(
|
|
collection: ResultCollection,
|
|
min_score: float,
|
|
*,
|
|
visible_bounds: tuple[float, float, float, float] | None = None,
|
|
object_draw_limits: tuple[int, int] | None = None,
|
|
) -> list[dict[str, float]]:
|
|
"""Build locator observations from GPR rows using score threshold and optional X/Z bounds."""
|
|
rows = gpr_object_rows(collection)
|
|
if rows.size == 0:
|
|
return []
|
|
|
|
finite_mask = np.all(np.isfinite(rows[:, :3]), axis=1)
|
|
visible_mask = finite_mask & (rows[:, 2] >= float(min_score))
|
|
if visible_bounds is not None:
|
|
x_min, x_max, z_min, z_max = (float(value) for value in visible_bounds)
|
|
visible_mask &= (
|
|
(rows[:, 0] >= x_min)
|
|
& (rows[:, 0] <= x_max)
|
|
& (rows[:, 1] >= z_min)
|
|
& (rows[:, 1] <= z_max)
|
|
)
|
|
filtered = rows[visible_mask]
|
|
if object_draw_limits is not None and filtered.size > 0:
|
|
max_detected_objects, draw_top_objects = object_draw_limits
|
|
if filtered.shape[0] > int(max_detected_objects):
|
|
filtered = np.zeros((0, filtered.shape[1]), dtype=filtered.dtype)
|
|
else:
|
|
filtered = filtered[: max(0, int(draw_top_objects))]
|
|
|
|
observations: list[dict[str, float]] = []
|
|
for x_m, z_m, _score in filtered:
|
|
observations.append(
|
|
{
|
|
"dst": round(float(z_m), 2),
|
|
"crs": round(float(x_m), 2),
|
|
}
|
|
)
|
|
return observations
|
|
|
|
|
|
def build_locator_payload(
|
|
observations: list[dict[str, float]],
|
|
*,
|
|
protocol_version: int,
|
|
status: int = 1,
|
|
) -> dict[str, Any]:
|
|
"""Assemble one outbound locator payload from precomputed observations."""
|
|
return {
|
|
"ver": int(protocol_version),
|
|
"tim": _format_timestamp(),
|
|
"sts": int(status),
|
|
"obs": observations,
|
|
}
|
|
|
|
|
|
def _format_timestamp() -> str:
|
|
"""Return wall-clock timestamp with millisecond precision."""
|
|
return datetime.now().strftime("%H:%M:%S.%f")[:-3]
|