1688 lines
58 KiB
Python
1688 lines
58 KiB
Python
"""Batch synthetic sonar dataset generator for PointNet semantic segmentation.
|
||
|
||
Produces paired Area_X_scene_XXXX.npy + .obj files under sonar_dataset/.
|
||
Each scene point is one echosounder return: for every beam×ping ray the first
|
||
hit (object or seafloor) is recorded; denser object returns come only from
|
||
sonar resolution and geometry, not from overlaying a second cloud.
|
||
Target class 1 = user-provided object; class 0 = seafloor.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import math
|
||
import random
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from scene_generator import (
|
||
apply_transform,
|
||
export_npy_float64,
|
||
export_obj,
|
||
generate_pipe,
|
||
parse_obj_labeled_points,
|
||
parse_obj_points,
|
||
)
|
||
|
||
# Full dataset layout (train / val / test).
|
||
AREA_LAYOUT: list[tuple[int, int]] = [
|
||
(1, 75),
|
||
(2, 75),
|
||
(3, 75),
|
||
(4, 75),
|
||
(5, 100),
|
||
(6, 100),
|
||
]
|
||
TOTAL_FULL_SCENES = sum(n for _, n in AREA_LAYOUT) # 500
|
||
|
||
VISIBILITY_TIERS = ("nearly_hidden", "partial", "visible")
|
||
|
||
# Built-in object models for the dataset generator (no upload required).
|
||
OBJECT_PRESETS: dict[str, dict[str, Any]] = {
|
||
"pipe": {
|
||
"id": "pipe",
|
||
"label": "Трубопровод",
|
||
"filename": "pipeline.obj",
|
||
# Length is scene-sized at placement time (border→border); params kept for catalog.
|
||
"params": {
|
||
"length": 2.6,
|
||
"radius": 0.22,
|
||
"axis": "y",
|
||
"count": 4000,
|
||
"noise": 0.005,
|
||
"seed": 1,
|
||
},
|
||
},
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Area naming
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def scene_index_to_area_name(index: int) -> tuple[int, int, str]:
|
||
"""Map 0-based global index → (area, scene_number_1based, stem).
|
||
|
||
Scene numbers restart at 0001 within each Area.
|
||
"""
|
||
if index < 0:
|
||
raise ValueError("scene index must be >= 0")
|
||
remaining = index
|
||
for area, count in AREA_LAYOUT:
|
||
if remaining < count:
|
||
scene_no = remaining + 1
|
||
stem = f"Area_{area}_scene_{scene_no:04d}"
|
||
return area, scene_no, stem
|
||
remaining -= count
|
||
scene_no = AREA_LAYOUT[-1][1] + remaining + 1
|
||
stem = f"Area_6_scene_{scene_no:04d}"
|
||
return 6, scene_no, stem
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Target object from user .obj or built-in preset
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def list_object_presets() -> list[dict[str, Any]]:
|
||
"""Public preset catalog for the dataset UI."""
|
||
return [
|
||
{
|
||
"id": meta["id"],
|
||
"label": meta["label"],
|
||
"filename": meta["filename"],
|
||
}
|
||
for meta in OBJECT_PRESETS.values()
|
||
]
|
||
|
||
|
||
def load_preset_object_points(preset_id: str) -> tuple[list[list[float]], str]:
|
||
"""Build normalized object points for a built-in preset. Returns (points, filename)."""
|
||
key = str(preset_id or "").strip().lower()
|
||
meta = OBJECT_PRESETS.get(key)
|
||
if meta is None:
|
||
known = ", ".join(sorted(OBJECT_PRESETS)) or "(none)"
|
||
raise ValueError(f"Unknown object preset '{preset_id}'. Known: {known}")
|
||
if key == "pipe":
|
||
points = generate_pipe(dict(meta["params"]))
|
||
else:
|
||
raise ValueError(f"Preset '{preset_id}' has no generator.")
|
||
if len(points) < 3:
|
||
raise ValueError(f"Preset '{preset_id}' produced too few points.")
|
||
return normalize_object_points(points), str(meta["filename"])
|
||
|
||
|
||
def normalize_object_points(points: list[list[float]]) -> list[list[float]]:
|
||
xs = [p[0] for p in points]
|
||
ys = [p[1] for p in points]
|
||
zs = [p[2] for p in points]
|
||
cx = (min(xs) + max(xs)) * 0.5
|
||
cy = (min(ys) + max(ys)) * 0.5
|
||
cz = (min(zs) + max(zs)) * 0.5
|
||
span = max(max(xs) - min(xs), max(ys) - min(ys), max(zs) - min(zs), 1e-6)
|
||
scale = 1.0 / span
|
||
return [[(p[0] - cx) * scale, (p[1] - cy) * scale, (p[2] - cz) * scale] for p in points]
|
||
|
||
|
||
def load_object_points_from_obj_text(text: str) -> list[list[float]]:
|
||
"""Parse OBJ vertices and normalize to unit local frame (centered, max span ≈ 1)."""
|
||
points = parse_obj_points(text)
|
||
if len(points) < 3:
|
||
raise ValueError("OBJ model must contain at least 3 vertices.")
|
||
return normalize_object_points(points)
|
||
|
||
|
||
def object_half_extent_z(points: list[list[float]]) -> float:
|
||
if not points:
|
||
return 0.35
|
||
zs = [p[2] for p in points]
|
||
return max(0.05, (max(zs) - min(zs)) * 0.5)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Seafloor heightfield + echosounder ray casting
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _seafloor_height(
|
||
x: float,
|
||
y: float,
|
||
*,
|
||
base_z: float,
|
||
amplitude: float,
|
||
frequency: float,
|
||
hills: list[tuple[float, float, float, float]],
|
||
valleys: list[tuple[float, float, float, float]],
|
||
bumps: list[tuple[float, float, float, float]],
|
||
) -> float:
|
||
z = base_z
|
||
z += amplitude * math.sin(frequency * x) * math.cos(frequency * 0.7 * y)
|
||
z += 0.35 * amplitude * math.sin(frequency * 1.7 * y + 0.4)
|
||
for cx, cy, height, radius in hills:
|
||
d2 = (x - cx) ** 2 + (y - cy) ** 2
|
||
if d2 < radius * radius * 4:
|
||
z += height * math.exp(-d2 / max(radius * radius, 1e-6))
|
||
for cx, cy, depth, radius in valleys:
|
||
d2 = (x - cx) ** 2 + (y - cy) ** 2
|
||
if d2 < radius * radius * 4:
|
||
z -= depth * math.exp(-d2 / max(radius * radius, 1e-6))
|
||
for cx, cy, height, radius in bumps:
|
||
d2 = (x - cx) ** 2 + (y - cy) ** 2
|
||
if d2 < radius * radius * 4:
|
||
z += height * math.exp(-d2 / max(radius * radius * 0.5, 1e-6))
|
||
return z
|
||
|
||
|
||
def _build_seafloor_meta(
|
||
rng: random.Random,
|
||
*,
|
||
beam_count: int = 45,
|
||
length_count: int | None = None,
|
||
mle_mode: bool = False,
|
||
auv_depth: float | None = None,
|
||
swath_angle_deg: float = 90.0,
|
||
relief_scale_pct: float = 100.0,
|
||
) -> dict[str, Any]:
|
||
"""Build continuous seafloor heightfield parameters (no point cloud yet).
|
||
|
||
``relief_scale_pct`` scales random unevenness relative to scene size
|
||
(100 = baseline, clearly visible relief; higher = larger features).
|
||
"""
|
||
beams = max(1, int(beam_count))
|
||
length_points = beams if length_count is None else max(1, int(length_count))
|
||
swath_deg = max(5.0, min(170.0, float(swath_angle_deg)))
|
||
relief_pct = max(5.0, min(500.0, float(relief_scale_pct)))
|
||
|
||
if mle_mode and auv_depth is not None:
|
||
# Swath footprint from altitude above flat datum (like МЛЭ).
|
||
depth = max(0.3, float(auv_depth))
|
||
half_w = depth * math.tan(math.radians(swath_deg) * 0.5)
|
||
size_x = max(4.0, 2.0 * half_w * 1.08)
|
||
dx = size_x / max(beams - 1, 1)
|
||
size_y = max(4.0, dx * max(length_points - 1, 1) * rng.uniform(0.95, 1.12))
|
||
else:
|
||
size_x = rng.uniform(8.0, 16.0)
|
||
size_y = rng.uniform(8.0, 16.0)
|
||
|
||
base_z = rng.uniform(-1.2, -0.2)
|
||
|
||
# Feature size tied to scene so relief stays visible in the point cloud.
|
||
# At 100%: characteristic scale ≈ 28% of the shorter scene side.
|
||
scene_ref = max(4.0, min(size_x, size_y))
|
||
feature_scale = scene_ref * (relief_pct / 100.0) * 0.28
|
||
feature_scale = max(0.35, feature_scale)
|
||
max_rad = scene_ref * 0.42
|
||
|
||
# Background undulation — several waves across the patch.
|
||
amplitude = feature_scale * rng.uniform(0.12, 0.28)
|
||
frequency = (2.0 * math.pi / scene_ref) * rng.uniform(1.8, 3.6)
|
||
|
||
n_hills = rng.randint(2, 5)
|
||
hills = []
|
||
for _ in range(n_hills):
|
||
hills.append(
|
||
(
|
||
rng.uniform(-size_x * 0.38, size_x * 0.38),
|
||
rng.uniform(-size_y * 0.38, size_y * 0.38),
|
||
feature_scale * rng.uniform(0.45, 1.05),
|
||
min(max_rad, feature_scale * rng.uniform(0.75, 1.55)),
|
||
)
|
||
)
|
||
|
||
n_valleys = rng.randint(1, 4)
|
||
valleys = []
|
||
for _ in range(n_valleys):
|
||
valleys.append(
|
||
(
|
||
rng.uniform(-size_x * 0.38, size_x * 0.38),
|
||
rng.uniform(-size_y * 0.38, size_y * 0.38),
|
||
feature_scale * rng.uniform(0.3, 0.75),
|
||
min(max_rad, feature_scale * rng.uniform(0.65, 1.4)),
|
||
)
|
||
)
|
||
|
||
# More medium bumps at lower scale; fewer but larger when scale is high.
|
||
density = max(0.45, min(1.8, 120.0 / max(relief_pct, 20.0)))
|
||
n_bumps = int(round(rng.uniform(10, 20) * density))
|
||
n_bumps = max(8, min(40, n_bumps))
|
||
bumps = []
|
||
for _ in range(n_bumps):
|
||
bumps.append(
|
||
(
|
||
rng.uniform(-size_x * 0.46, size_x * 0.46),
|
||
rng.uniform(-size_y * 0.46, size_y * 0.46),
|
||
feature_scale * rng.uniform(0.12, 0.45),
|
||
min(max_rad * 0.7, feature_scale * rng.uniform(0.25, 0.85)),
|
||
)
|
||
)
|
||
|
||
return {
|
||
"sizeX": size_x,
|
||
"sizeY": size_y,
|
||
"baseZ": base_z,
|
||
"amplitude": amplitude,
|
||
"frequency": frequency,
|
||
"hills": hills,
|
||
"valleys": valleys,
|
||
"bumps": bumps,
|
||
"beamCount": beams,
|
||
"lengthCount": length_points,
|
||
"gridWidthPoints": beams,
|
||
"gridLengthPoints": length_points,
|
||
"pingCount": length_points,
|
||
"swathBeams": beams,
|
||
"noise": rng.uniform(0.002, 0.02),
|
||
"mleMode": bool(mle_mode),
|
||
"auvDepth": float(auv_depth) if auv_depth is not None else None,
|
||
"swathAngleDeg": swath_deg,
|
||
"reliefScalePct": relief_pct,
|
||
"featureScale": feature_scale,
|
||
}
|
||
|
||
|
||
def _height_at(x: float, y: float, meta: dict[str, Any]) -> float:
|
||
return _seafloor_height(
|
||
x,
|
||
y,
|
||
base_z=float(meta["baseZ"]),
|
||
amplitude=float(meta["amplitude"]),
|
||
frequency=float(meta["frequency"]),
|
||
hills=meta["hills"],
|
||
valleys=meta["valleys"],
|
||
bumps=meta["bumps"],
|
||
)
|
||
|
||
|
||
def _grid_xy(xi: int, yi: int, meta: dict[str, Any]) -> tuple[float, float]:
|
||
beams = int(meta["beamCount"])
|
||
length_points = int(meta["lengthCount"])
|
||
size_x = float(meta["sizeX"])
|
||
size_y = float(meta["sizeY"])
|
||
half_x = size_x * 0.5
|
||
half_y = size_y * 0.5
|
||
x = -half_x if beams == 1 else (-half_x + size_x * xi / (beams - 1))
|
||
y = -half_y if length_points == 1 else (-half_y + size_y * yi / (length_points - 1))
|
||
return x, y
|
||
|
||
|
||
def _cell_size(meta: dict[str, Any]) -> tuple[float, float]:
|
||
beams = int(meta["beamCount"])
|
||
length_points = int(meta["lengthCount"])
|
||
size_x = float(meta["sizeX"])
|
||
size_y = float(meta["sizeY"])
|
||
return size_x / max(beams - 1, 1), size_y / max(length_points - 1, 1)
|
||
|
||
|
||
def _world_to_grid_index(
|
||
x: float,
|
||
y: float,
|
||
meta: dict[str, Any],
|
||
) -> tuple[float, float]:
|
||
beams = int(meta["beamCount"])
|
||
length_points = int(meta["lengthCount"])
|
||
size_x = float(meta["sizeX"])
|
||
size_y = float(meta["sizeY"])
|
||
half_x = size_x * 0.5
|
||
half_y = size_y * 0.5
|
||
fx = 0.0 if beams == 1 else (x + half_x) / size_x * (beams - 1)
|
||
fy = 0.0 if length_points == 1 else (y + half_y) / size_y * (length_points - 1)
|
||
return fx, fy
|
||
|
||
|
||
def _place_object_in_scene(
|
||
rng: random.Random,
|
||
meta: dict[str, Any],
|
||
visibility: str,
|
||
object_template: list[list[float]],
|
||
object_scale: float = 1.0,
|
||
) -> tuple[list[list[float]], dict[str, Any]]:
|
||
"""Pose the object on the seafloor; return world-space template vertices + info."""
|
||
base_scale = max(0.01, float(object_scale))
|
||
if visibility == "nearly_hidden":
|
||
burial = rng.uniform(0.35, 0.75)
|
||
scale = base_scale * rng.uniform(0.7, 1.15)
|
||
elif visibility == "partial":
|
||
burial = rng.uniform(0.12, 0.4)
|
||
scale = base_scale * rng.uniform(0.8, 1.3)
|
||
else:
|
||
burial = rng.uniform(-0.05, 0.15)
|
||
scale = base_scale * rng.uniform(0.85, 1.4)
|
||
|
||
local = [[p[0] * scale, p[1] * scale, p[2] * scale] for p in object_template]
|
||
half_x = float(meta["sizeX"]) * 0.35
|
||
half_y = float(meta["sizeY"]) * 0.35
|
||
tx = rng.uniform(-half_x, half_x)
|
||
ty = rng.uniform(-half_y, half_y)
|
||
floor_z = _height_at(tx, ty, meta)
|
||
half_h = object_half_extent_z(local)
|
||
tz = floor_z + half_h * (1.0 - 2.0 * burial)
|
||
transform = {
|
||
"x": tx,
|
||
"y": ty,
|
||
"z": tz,
|
||
"rx": rng.uniform(-0.25, 0.25),
|
||
"ry": rng.uniform(-0.2, 0.2),
|
||
"rz": rng.uniform(0, 2 * math.pi),
|
||
}
|
||
world = apply_transform(local, transform)
|
||
info = {
|
||
"visibility": visibility,
|
||
"transform": transform,
|
||
"scale": scale,
|
||
"objectScale": base_scale,
|
||
"burial": burial,
|
||
"vertexCount": len(world),
|
||
"classLabel": "object",
|
||
"classId": 1,
|
||
}
|
||
return world, info
|
||
|
||
|
||
def _terrain_unevenness(meta: dict[str, Any]) -> float:
|
||
"""Summarize seafloor relief strength as a ~0…1 scalar for silt scaling."""
|
||
amp = float(meta.get("amplitude", 0.0))
|
||
hills = meta.get("hills") or []
|
||
valleys = meta.get("valleys") or []
|
||
bumps = meta.get("bumps") or []
|
||
hill_peak = max((float(h[2]) for h in hills), default=0.0)
|
||
valley_peak = max((float(v[2]) for v in valleys), default=0.0)
|
||
bump_peak = max((float(b[2]) for b in bumps), default=0.0)
|
||
score = (
|
||
amp / 0.35
|
||
+ 0.5 * (hill_peak / 0.7)
|
||
+ 0.3 * (valley_peak / 0.55)
|
||
+ 0.25 * (bump_peak / 0.35)
|
||
+ 0.04 * len(hills)
|
||
+ 0.02 * len(bumps)
|
||
)
|
||
return max(0.0, min(1.0, score / 2.2))
|
||
|
||
|
||
def _place_pipeline_in_scene(
|
||
rng: random.Random,
|
||
meta: dict[str, Any],
|
||
visibility: str,
|
||
object_scale: float = 1.0,
|
||
) -> tuple[list[list[float]], dict[str, Any]]:
|
||
"""Lay a straight pipeline from scan border to border; may be silted over.
|
||
|
||
Length spans the full survey rectangle along X or Y (small yaw allowed).
|
||
For ``visible`` there is no silt. For partial / nearly_hidden, randomly
|
||
one silt mound or several (count from visibility and seafloor unevenness).
|
||
"""
|
||
size_x = float(meta["sizeX"])
|
||
size_y = float(meta["sizeY"])
|
||
min_span = min(size_x, size_y)
|
||
base_scale = max(0.01, float(object_scale))
|
||
uneven = _terrain_unevenness(meta)
|
||
|
||
# Diameter scales with scene and UI objectScale; length always edge-to-edge.
|
||
radius = max(0.05, min_span * 0.028 * base_scale)
|
||
radius = min(radius, min_span * 0.12)
|
||
|
||
along_x = rng.random() < 0.5
|
||
yaw = rng.uniform(-math.radians(10), math.radians(10))
|
||
cos_y = math.cos(yaw)
|
||
sin_y = math.sin(yaw)
|
||
|
||
if along_x:
|
||
length = size_x / max(abs(cos_y), 0.88) * 1.04
|
||
cx = 0.0
|
||
cy = rng.uniform(-size_y * 0.12, size_y * 0.12)
|
||
ux, uy = cos_y, sin_y
|
||
nx, ny = -sin_y, cos_y
|
||
else:
|
||
length = size_y / max(abs(cos_y), 0.88) * 1.04
|
||
cx = rng.uniform(-size_x * 0.12, size_x * 0.12)
|
||
cy = 0.0
|
||
ux, uy = -sin_y, cos_y
|
||
nx, ny = -cos_y, -sin_y
|
||
|
||
def centerline_xy(t_norm: float) -> tuple[float, float]:
|
||
t = t_norm * length
|
||
return cx + ux * t, cy + uy * t
|
||
|
||
# Ось — прямая в 3D: одна высота для всего цилиндра (не повторяет рельеф).
|
||
n_along = max(96, int(max(int(meta["beamCount"]), int(meta["lengthCount"])) * 3))
|
||
n_circ = 28
|
||
floor_samples = [
|
||
_height_at(*centerline_xy(ia / max(n_along - 1, 1) - 0.5), meta)
|
||
for ia in range(n_along)
|
||
]
|
||
# Положить трубу примерно на средний уровень дна, не изгибая ось.
|
||
axis_z = (sum(floor_samples) / max(len(floor_samples), 1)) + radius
|
||
|
||
world: list[list[float]] = []
|
||
for ia in range(n_along):
|
||
t_norm = ia / max(n_along - 1, 1) - 0.5 # [-0.5, 0.5]
|
||
px, py = centerline_xy(t_norm)
|
||
for ic in range(n_circ):
|
||
ang = (2.0 * math.pi * ic) / n_circ
|
||
rr = math.cos(ang) * radius
|
||
rz = math.sin(ang) * radius
|
||
world.append([px + nx * rr, py + ny * rr, axis_z + rz])
|
||
|
||
# Ил: только partial / nearly_hidden — случайно один холм или несколько.
|
||
silt_bumps: list[tuple[float, float, float, float]] = []
|
||
silt_mode = "none"
|
||
if visibility in ("nearly_hidden", "partial"):
|
||
silt_mode = "several" if rng.random() < 0.5 else "one"
|
||
if silt_mode == "one":
|
||
n_mounds = 1
|
||
elif visibility == "nearly_hidden":
|
||
n_mounds = rng.randint(2, 5) + (1 if uneven > 0.55 else 0)
|
||
else:
|
||
n_mounds = rng.randint(2, 4) + (1 if uneven > 0.65 else 0)
|
||
|
||
if visibility == "nearly_hidden":
|
||
span = 0.22 + 0.28 * uneven
|
||
cover = 0.85 + 0.55 * uneven
|
||
else:
|
||
span = 0.14 + 0.22 * uneven
|
||
cover = 0.4 + 0.45 * uneven
|
||
|
||
# Один холм — шире/выше; несколько — компактнее и разнесены.
|
||
for i in range(n_mounds):
|
||
if n_mounds == 1:
|
||
t_norm = rng.uniform(-span * 0.55, span * 0.55)
|
||
else:
|
||
t_norm = -span + (2.0 * span) * (i + 0.5) / n_mounds
|
||
t_norm += rng.uniform(-span * 0.12, span * 0.12)
|
||
t_norm = max(-span, min(span, t_norm))
|
||
lateral = rng.uniform(-radius * (1.0 + uneven), radius * (1.0 + uneven))
|
||
px, py = centerline_xy(t_norm)
|
||
sx = px + nx * lateral
|
||
sy = py + ny * lateral
|
||
local_floor = _height_at(sx, sy, meta)
|
||
neighbor = _height_at(sx + radius * 2, sy + radius * 2, meta)
|
||
local_relief = abs(local_floor - neighbor) / max(radius, 1e-3)
|
||
local_factor = 0.75 + min(0.55, 0.35 * local_relief + 0.25 * uneven)
|
||
size_boost = 1.35 if n_mounds == 1 else 1.0
|
||
height = radius * cover * local_factor * size_boost * rng.uniform(0.85, 1.2)
|
||
mound_r = max(
|
||
radius * rng.uniform(2.8, 4.8) * (1.35 if n_mounds == 1 else 1.0),
|
||
min_span * rng.uniform(0.05, 0.1) * (0.85 + 0.3 * uneven),
|
||
)
|
||
silt_bumps.append((sx, sy, height, mound_r))
|
||
else:
|
||
cover = 0.0
|
||
|
||
if silt_bumps:
|
||
meta.setdefault("bumps", []).extend(silt_bumps)
|
||
|
||
info = {
|
||
"visibility": visibility,
|
||
"kind": "pipe",
|
||
"transform": {
|
||
"x": cx,
|
||
"y": cy,
|
||
"z": 0.0,
|
||
"rx": 0.0,
|
||
"ry": 0.0,
|
||
"rz": yaw if along_x else (yaw + 0.5 * math.pi),
|
||
},
|
||
"scale": base_scale,
|
||
"objectScale": base_scale,
|
||
"radius": radius,
|
||
"length": length,
|
||
"axis": "x" if along_x else "y",
|
||
"hasBend": False,
|
||
"bendAmp": 0.0,
|
||
"axisZ": round(axis_z, 4),
|
||
"burial": 0.0,
|
||
"terrainUnevenness": round(uneven, 4),
|
||
"siltMode": silt_mode,
|
||
"siltCover": round(cover, 4),
|
||
"siltBumpCount": len(silt_bumps),
|
||
"vertexCount": len(world),
|
||
"classLabel": "object",
|
||
"classId": 1,
|
||
}
|
||
return world, info
|
||
|
||
|
||
def _rasterize_object_hits(
|
||
world_pts: list[list[float]],
|
||
meta: dict[str, Any],
|
||
) -> dict[tuple[int, int], float]:
|
||
"""Project object vertices onto the sonar grid: first-hit Z per beam×ping cell.
|
||
|
||
Sensor looks down (+Z is closer). A cell stores the highest object Z that
|
||
falls into its footprint, so later casting can compare against seafloor Z.
|
||
"""
|
||
beams = int(meta["beamCount"])
|
||
length_points = int(meta["lengthCount"])
|
||
cell_w, cell_h = _cell_size(meta)
|
||
# Cover gaps between sparse mesh vertices across neighboring cells
|
||
splat_rx = cell_w * 0.85
|
||
splat_ry = cell_h * 0.85
|
||
hits: dict[tuple[int, int], float] = {}
|
||
|
||
for px, py, pz in world_pts:
|
||
fx, fy = _world_to_grid_index(px, py, meta)
|
||
xi0 = int(math.floor(fx))
|
||
yi0 = int(math.floor(fy))
|
||
for dyi in (-1, 0, 1, 2):
|
||
for dxi in (-1, 0, 1, 2):
|
||
xi = xi0 + dxi
|
||
yi = yi0 + dyi
|
||
if xi < 0 or yi < 0 or xi >= beams or yi >= length_points:
|
||
continue
|
||
cx, cy = _grid_xy(xi, yi, meta)
|
||
if abs(px - cx) > splat_rx or abs(py - cy) > splat_ry:
|
||
continue
|
||
floor_z = _height_at(cx, cy, meta)
|
||
# Buried volume does not return a sonar echo above the seafloor
|
||
if pz < floor_z - 0.01:
|
||
continue
|
||
key = (xi, yi)
|
||
prev = hits.get(key)
|
||
if prev is None or pz > prev:
|
||
hits[key] = pz
|
||
return hits
|
||
|
||
|
||
def _cast_echosounder_returns(
|
||
rng: random.Random,
|
||
meta: dict[str, Any],
|
||
object_hits: dict[tuple[int, int], float] | None,
|
||
) -> tuple[list[list[float]], int, int]:
|
||
"""One return per ray: first surface hit from above (object or seafloor).
|
||
|
||
Returns PointNet rows [x,y,z,r,g,b,class], object_count, background_count.
|
||
"""
|
||
beams = int(meta["beamCount"])
|
||
length_points = int(meta["lengthCount"])
|
||
noise = float(meta.get("noise", 0.01))
|
||
rows: list[list[float]] = []
|
||
object_count = 0
|
||
background_count = 0
|
||
|
||
for yi in range(length_points):
|
||
for xi in range(beams):
|
||
x, y = _grid_xy(xi, yi, meta)
|
||
floor_z = _height_at(x, y, meta)
|
||
obj_z = object_hits.get((xi, yi)) if object_hits else None
|
||
|
||
# Looking down: larger Z is closer → first intersection wins.
|
||
if obj_z is not None and obj_z > floor_z:
|
||
z = obj_z + rng.uniform(-noise, noise)
|
||
cls = 1.0
|
||
object_count += 1
|
||
else:
|
||
z = floor_z + rng.uniform(-noise, noise)
|
||
cls = 0.0
|
||
background_count += 1
|
||
|
||
rows.append([float(x), float(y), float(z), 0.0, 0.0, 0.0, cls])
|
||
|
||
return rows, object_count, background_count
|
||
|
||
|
||
def _build_object_height_field(
|
||
world_pts: list[list[float]],
|
||
meta: dict[str, Any],
|
||
) -> dict[str, Any] | None:
|
||
"""Dense XY max-Z map of object surface for angled multibeam probes."""
|
||
if not world_pts:
|
||
return None
|
||
beams = int(meta["beamCount"])
|
||
length_points = int(meta["lengthCount"])
|
||
nx = max(32, min(256, beams * 3))
|
||
ny = max(32, min(256, length_points * 3))
|
||
size_x = float(meta["sizeX"])
|
||
size_y = float(meta["sizeY"])
|
||
half_x = size_x * 0.5
|
||
half_y = size_y * 0.5
|
||
cell_w = size_x / max(nx - 1, 1)
|
||
cell_h = size_y / max(ny - 1, 1)
|
||
splat_rx = cell_w * 1.1
|
||
splat_ry = cell_h * 1.1
|
||
field: list[list[float | None]] = [[None for _ in range(nx)] for _ in range(ny)]
|
||
|
||
for px, py, pz in world_pts:
|
||
fx = 0.0 if nx == 1 else (px + half_x) / size_x * (nx - 1)
|
||
fy = 0.0 if ny == 1 else (py + half_y) / size_y * (ny - 1)
|
||
xi0 = int(math.floor(fx))
|
||
yi0 = int(math.floor(fy))
|
||
for dyi in (-1, 0, 1, 2):
|
||
for dxi in (-1, 0, 1, 2):
|
||
xi = xi0 + dxi
|
||
yi = yi0 + dyi
|
||
if xi < 0 or yi < 0 or xi >= nx or yi >= ny:
|
||
continue
|
||
cx = -half_x if nx == 1 else (-half_x + size_x * xi / (nx - 1))
|
||
cy = -half_y if ny == 1 else (-half_y + size_y * yi / (ny - 1))
|
||
if abs(px - cx) > splat_rx or abs(py - cy) > splat_ry:
|
||
continue
|
||
floor_z = _height_at(cx, cy, meta)
|
||
if pz < floor_z - 0.01:
|
||
continue
|
||
prev = field[yi][xi]
|
||
if prev is None or pz > prev:
|
||
field[yi][xi] = pz
|
||
|
||
return {
|
||
"nx": nx,
|
||
"ny": ny,
|
||
"sizeX": size_x,
|
||
"sizeY": size_y,
|
||
"field": field,
|
||
}
|
||
|
||
|
||
def _sample_object_height_field(
|
||
field_info: dict[str, Any] | None,
|
||
x: float,
|
||
y: float,
|
||
) -> float | None:
|
||
if not field_info:
|
||
return None
|
||
nx = int(field_info["nx"])
|
||
ny = int(field_info["ny"])
|
||
size_x = float(field_info["sizeX"])
|
||
size_y = float(field_info["sizeY"])
|
||
half_x = size_x * 0.5
|
||
half_y = size_y * 0.5
|
||
u = (x + half_x) / max(size_x, 1e-6)
|
||
v = (y + half_y) / max(size_y, 1e-6)
|
||
if u < -0.02 or u > 1.02 or v < -0.02 or v > 1.02:
|
||
return None
|
||
u = min(1.0, max(0.0, u))
|
||
v = min(1.0, max(0.0, v))
|
||
fx = u * (nx - 1)
|
||
fy = v * (ny - 1)
|
||
i0 = int(math.floor(fx))
|
||
j0 = int(math.floor(fy))
|
||
i1 = min(nx - 1, i0 + 1)
|
||
j1 = min(ny - 1, j0 + 1)
|
||
vals = [
|
||
field_info["field"][j0][i0],
|
||
field_info["field"][j0][i1],
|
||
field_info["field"][j1][i0],
|
||
field_info["field"][j1][i1],
|
||
]
|
||
present = [z for z in vals if z is not None]
|
||
if not present:
|
||
return None
|
||
return max(present)
|
||
|
||
|
||
def _cast_mle_fan_returns(
|
||
rng: random.Random,
|
||
meta: dict[str, Any],
|
||
object_field: dict[str, Any] | None,
|
||
*,
|
||
auv_depth: float,
|
||
swath_angle_deg: float = 90.0,
|
||
) -> tuple[list[list[float]], int, int]:
|
||
"""Multibeam fan: beamCount rays × lengthCount pings from constant altitude.
|
||
|
||
AUV height is measured from flat seafloor datum (baseZ), not terrain-following.
|
||
"""
|
||
beams = int(meta["beamCount"])
|
||
length_points = int(meta["lengthCount"])
|
||
noise = float(meta.get("noise", 0.01))
|
||
base_z = float(meta["baseZ"])
|
||
depth = max(0.3, float(auv_depth))
|
||
origin_z = base_z + depth
|
||
swath = math.radians(max(5.0, min(170.0, float(swath_angle_deg))))
|
||
size_y = float(meta["sizeY"])
|
||
half_y = size_y * 0.5
|
||
# Reach outer beams with margin past the swath edge.
|
||
max_range = max(depth / max(math.cos(swath * 0.5), 0.12) * 1.35, depth * 2.0, 4.0)
|
||
step = max(0.08, min(0.45, max_range / 220.0))
|
||
max_steps = int(math.ceil(max_range / step)) + 2
|
||
|
||
rows: list[list[float]] = []
|
||
object_count = 0
|
||
background_count = 0
|
||
|
||
for yi in range(length_points):
|
||
y = -half_y if length_points == 1 else (-half_y + size_y * yi / (length_points - 1))
|
||
ox, oy, oz = 0.0, y, origin_z
|
||
for xi in range(beams):
|
||
t = 0.5 if beams == 1 else xi / (beams - 1)
|
||
angle = -swath * 0.5 + swath * t
|
||
dx = math.sin(angle)
|
||
dy = 0.0
|
||
dz = -math.cos(angle)
|
||
|
||
hit_x, hit_y, hit_z = ox, oy, base_z
|
||
is_object = False
|
||
px = ox + dx * step * 0.5
|
||
py = oy + dy * step * 0.5
|
||
pz = oz + dz * step * 0.5
|
||
for _ in range(max_steps):
|
||
dist = math.hypot(px - ox, py - oy, pz - oz)
|
||
if dist > max_range:
|
||
break
|
||
floor_z = _height_at(px, py, meta)
|
||
obj_z = _sample_object_height_field(object_field, px, py)
|
||
if obj_z is not None and pz <= obj_z and obj_z >= floor_z - 0.01:
|
||
hit_x, hit_y, hit_z = px, py, obj_z
|
||
is_object = True
|
||
break
|
||
if pz <= floor_z:
|
||
hit_x, hit_y, hit_z = px, py, floor_z
|
||
break
|
||
px += dx * step
|
||
py += dy * step
|
||
pz += dz * step
|
||
else:
|
||
# No intersection within range — project to flat nadir estimate.
|
||
hit_x = ox + dx * max_range
|
||
hit_y = oy + dy * max_range
|
||
hit_z = _height_at(hit_x, hit_y, meta)
|
||
|
||
z = hit_z + rng.uniform(-noise, noise)
|
||
if is_object:
|
||
cls = 1.0
|
||
object_count += 1
|
||
else:
|
||
cls = 0.0
|
||
background_count += 1
|
||
rows.append([float(hit_x), float(hit_y), float(z), 0.0, 0.0, 0.0, cls])
|
||
|
||
return rows, object_count, background_count
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Balance plan + single scene
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _allocate_by_weights(total: int, weights: list[float]) -> list[int]:
|
||
"""Split ``total`` into integer buckets proportional to non-negative weights."""
|
||
n = len(weights)
|
||
if total <= 0 or n == 0:
|
||
return [0] * n
|
||
w = [max(0.0, float(x)) for x in weights]
|
||
s = sum(w)
|
||
if s <= 0:
|
||
out = [0] * n
|
||
out[0] = total
|
||
return out
|
||
raw = [total * wi / s for wi in w]
|
||
floors = [int(math.floor(r)) for r in raw]
|
||
rem = total - sum(floors)
|
||
order = sorted(range(n), key=lambda i: (raw[i] - floors[i], -i), reverse=True)
|
||
for i in range(rem):
|
||
floors[order[i % n]] += 1
|
||
return floors
|
||
|
||
|
||
def _clamp_pct(value: float, *, name: str) -> float:
|
||
v = float(value)
|
||
if not math.isfinite(v):
|
||
raise ValueError(f"{name} must be a finite number")
|
||
if v < 0 or v > 100:
|
||
raise ValueError(f"{name} must be in [0, 100], got {v}")
|
||
return v
|
||
|
||
|
||
def plan_scene_labels(
|
||
count: int,
|
||
seed: int,
|
||
*,
|
||
absent_pct: float = 30.0,
|
||
nearly_hidden_pct: float = 20.0,
|
||
partial_pct: float = 40.0,
|
||
visible_pct: float = 40.0,
|
||
) -> list[str]:
|
||
"""Return visibility label per scene: absent | nearly_hidden | partial | visible.
|
||
|
||
Defaults: 30% absent / 70% with object. Of scenes with an object:
|
||
20% nearly_hidden, 40% partial, 40% visible.
|
||
Visibility percents are relative to with-object scenes (normalized if needed).
|
||
"""
|
||
count = max(0, int(count))
|
||
absent_pct = _clamp_pct(absent_pct, name="absent_pct")
|
||
nearly_hidden_pct = _clamp_pct(nearly_hidden_pct, name="nearly_hidden_pct")
|
||
partial_pct = _clamp_pct(partial_pct, name="partial_pct")
|
||
visible_pct = _clamp_pct(visible_pct, name="visible_pct")
|
||
|
||
n_without, n_with = _allocate_by_weights(count, [absent_pct, 100.0 - absent_pct])
|
||
tier_weights = [nearly_hidden_pct, partial_pct, visible_pct]
|
||
if n_with > 0 and sum(tier_weights) <= 0:
|
||
raise ValueError(
|
||
"Visibility percents among with-object scenes must sum to > 0 "
|
||
"when there are scenes with an object."
|
||
)
|
||
n_nearly, n_partial, n_visible = _allocate_by_weights(n_with, tier_weights)
|
||
|
||
labels: list[str] = (
|
||
["absent"] * n_without
|
||
+ ["nearly_hidden"] * n_nearly
|
||
+ ["partial"] * n_partial
|
||
+ ["visible"] * n_visible
|
||
)
|
||
rng = random.Random(int(seed) ^ 0xA5A5_5A5A)
|
||
rng.shuffle(labels)
|
||
return labels
|
||
|
||
|
||
def generate_sonar_scene(
|
||
*,
|
||
seed: int,
|
||
visibility: str = "absent",
|
||
object_points: list[list[float]] | None = None,
|
||
object_scale: float = 1.0,
|
||
object_kind: str | None = None,
|
||
beam_count: int = 45,
|
||
length_count: int | None = None,
|
||
mle_mode: bool = False,
|
||
auv_depth: float | None = None,
|
||
swath_angle_deg: float = 90.0,
|
||
relief_scale_pct: float = 100.0,
|
||
) -> dict[str, Any]:
|
||
"""Build one sonar scene via beam×ping first-hit casting.
|
||
|
||
visibility in absent|nearly_hidden|partial|visible.
|
||
Scene size is always beam_count × length_count returns.
|
||
object_kind ``pipe`` places a border-to-border pipeline with optional silt.
|
||
|
||
When ``mle_mode`` is True, returns are cast as a multibeam fan from constant
|
||
altitude ``auv_depth`` above the flat seafloor datum (baseZ).
|
||
"""
|
||
rng = random.Random(int(seed))
|
||
if visibility not in ("absent",) + VISIBILITY_TIERS:
|
||
raise ValueError(f"Unknown visibility: {visibility}")
|
||
kind = (object_kind or "").strip().lower() or None
|
||
if visibility != "absent" and kind != "pipe" and not object_points:
|
||
raise ValueError("object_points required when visibility is not absent.")
|
||
|
||
use_mle = bool(mle_mode)
|
||
depth = max(0.3, float(auv_depth)) if auv_depth is not None else None
|
||
if use_mle and depth is None:
|
||
depth = 2.5
|
||
|
||
meta = _build_seafloor_meta(
|
||
rng,
|
||
beam_count=beam_count,
|
||
length_count=length_count,
|
||
mle_mode=use_mle,
|
||
auv_depth=depth,
|
||
swath_angle_deg=swath_angle_deg,
|
||
relief_scale_pct=relief_scale_pct,
|
||
)
|
||
expected = int(meta["beamCount"]) * int(meta["lengthCount"])
|
||
|
||
object_info: dict[str, Any] | None = None
|
||
object_hits: dict[tuple[int, int], float] | None = None
|
||
object_field: dict[str, Any] | None = None
|
||
world: list[list[float]] | None = None
|
||
if visibility != "absent":
|
||
if kind == "pipe":
|
||
world, object_info = _place_pipeline_in_scene(
|
||
rng,
|
||
meta,
|
||
visibility,
|
||
object_scale=object_scale,
|
||
)
|
||
else:
|
||
world, object_info = _place_object_in_scene(
|
||
rng,
|
||
meta,
|
||
visibility,
|
||
object_points,
|
||
object_scale=object_scale,
|
||
)
|
||
if use_mle:
|
||
object_field = _build_object_height_field(world, meta)
|
||
hit_cells = 0
|
||
if object_field:
|
||
for row in object_field["field"]:
|
||
hit_cells += sum(1 for z in row if z is not None)
|
||
object_info["hitCellCount"] = hit_cells
|
||
object_info["keptCount"] = hit_cells
|
||
else:
|
||
object_hits = _rasterize_object_hits(world, meta)
|
||
object_info["hitCellCount"] = len(object_hits)
|
||
object_info["keptCount"] = len(object_hits)
|
||
object_info["requestedCount"] = expected
|
||
|
||
if use_mle:
|
||
rows, object_point_count, background_point_count = _cast_mle_fan_returns(
|
||
rng,
|
||
meta,
|
||
object_field,
|
||
auv_depth=float(depth),
|
||
swath_angle_deg=float(meta.get("swathAngleDeg", swath_angle_deg)),
|
||
)
|
||
else:
|
||
rows, object_point_count, background_point_count = _cast_echosounder_returns(
|
||
rng, meta, object_hits
|
||
)
|
||
if len(rows) != expected:
|
||
raise RuntimeError(f"Ray count mismatch: expected {expected}, got {len(rows)}")
|
||
|
||
rng.shuffle(rows)
|
||
xyz = [[r[0], r[1], r[2]] for r in rows]
|
||
return {
|
||
"seed": int(seed),
|
||
"visibility": visibility,
|
||
"hasObject": visibility != "absent",
|
||
"object": object_info,
|
||
"pointCount": len(rows),
|
||
"objectPointCount": object_point_count,
|
||
"backgroundPointCount": background_point_count,
|
||
"auvDepth": float(depth) if use_mle else None,
|
||
"mleMode": use_mle,
|
||
"rows": rows,
|
||
"points": xyz,
|
||
"meta": {
|
||
"sizeX": meta["sizeX"],
|
||
"sizeY": meta["sizeY"],
|
||
"beamCount": meta.get("beamCount", beam_count),
|
||
"lengthCount": meta.get(
|
||
"lengthCount",
|
||
length_count if length_count is not None else beam_count,
|
||
),
|
||
"gridWidthPoints": meta.get("gridWidthPoints", beam_count),
|
||
"gridLengthPoints": meta.get(
|
||
"gridLengthPoints",
|
||
length_count if length_count is not None else beam_count,
|
||
),
|
||
"pingCount": meta.get("pingCount"),
|
||
"swathBeams": meta.get("swathBeams"),
|
||
"gridPointCount": expected,
|
||
"mleMode": use_mle,
|
||
"auvDepth": float(depth) if use_mle else None,
|
||
"swathAngleDeg": meta.get("swathAngleDeg"),
|
||
"reliefScalePct": meta.get("reliefScalePct"),
|
||
"floorFeatures": {
|
||
"hills": len(meta["hills"]),
|
||
"valleys": len(meta["valleys"]),
|
||
"bumps": len(meta["bumps"]),
|
||
},
|
||
},
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Batch write / preview load
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def resolve_output_dir(output_dir: str | Path = "sonar_dataset") -> Path:
|
||
out = Path(output_dir)
|
||
if not out.is_absolute():
|
||
project_root = Path(__file__).resolve().parent.parent
|
||
out = project_root / out
|
||
return out
|
||
|
||
|
||
DATASET_RUN_FILENAME = "dataset_run.json"
|
||
DATASET_RUN_VERSION = 1
|
||
|
||
|
||
def _safe_object_stem(object_name: str | None) -> str:
|
||
stem = Path(object_name or "object").stem
|
||
safe = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in stem).strip("._-")
|
||
return (safe or "object")[:80]
|
||
|
||
|
||
def make_generation_run_dir(base_dir: Path, object_name: str | None = None) -> Path:
|
||
"""Create a new run subdirectory: YYYY-MM-DD_HH-MM-SS-<object_stem>.
|
||
|
||
Previous runs under base_dir are left untouched.
|
||
"""
|
||
from datetime import datetime
|
||
|
||
base_dir.mkdir(parents=True, exist_ok=True)
|
||
stamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||
folder = f"{stamp}-{_safe_object_stem(object_name)}"
|
||
path = base_dir / folder
|
||
if path.exists():
|
||
n = 2
|
||
while True:
|
||
candidate = base_dir / f"{folder}_{n}"
|
||
if not candidate.exists():
|
||
path = candidate
|
||
break
|
||
n += 1
|
||
path.mkdir(parents=True, exist_ok=False)
|
||
return path
|
||
|
||
|
||
def _normalize_model_filename(name: str | None) -> str | None:
|
||
"""Keep the full uploaded basename, e.g. ``airplane2.obj``."""
|
||
if not name:
|
||
return None
|
||
base = Path(str(name).strip()).name
|
||
return base or None
|
||
|
||
|
||
def build_run_manifest(
|
||
*,
|
||
run_dir: Path,
|
||
base_dir: Path,
|
||
count: int,
|
||
seed: int,
|
||
output_dir: str,
|
||
object_name: str | None,
|
||
object_scale: float,
|
||
object_scale_is_max: bool,
|
||
beam_count: int,
|
||
length_count: int,
|
||
object_vertex_count: int,
|
||
stats: dict[str, Any],
|
||
written: list[dict[str, Any]],
|
||
absent_pct: float = 30.0,
|
||
nearly_hidden_pct: float = 20.0,
|
||
partial_pct: float = 40.0,
|
||
visible_pct: float = 40.0,
|
||
mle_mode: bool = False,
|
||
auv_depth_min: float | None = None,
|
||
auv_depth_max: float | None = None,
|
||
relief_scale_pct: float = 100.0,
|
||
) -> dict[str, Any]:
|
||
model_filename = _normalize_model_filename(object_name)
|
||
settings: dict[str, Any] = {
|
||
"count": int(count),
|
||
"seed": int(seed),
|
||
"outputDir": str(output_dir),
|
||
"objectScale": float(object_scale),
|
||
"objectScaleIsMax": bool(object_scale_is_max),
|
||
"beamCount": int(beam_count),
|
||
"lengthCount": int(length_count),
|
||
"objectName": model_filename,
|
||
"absentPct": float(absent_pct),
|
||
"nearlyHiddenPct": float(nearly_hidden_pct),
|
||
"partialPct": float(partial_pct),
|
||
"visiblePct": float(visible_pct),
|
||
"mleMode": bool(mle_mode),
|
||
"reliefScalePct": float(relief_scale_pct),
|
||
}
|
||
if mle_mode:
|
||
settings["auvDepthMin"] = float(auv_depth_min) if auv_depth_min is not None else None
|
||
settings["auvDepthMax"] = float(auv_depth_max) if auv_depth_max is not None else None
|
||
return {
|
||
"version": DATASET_RUN_VERSION,
|
||
"generatedAt": datetime.now(timezone.utc).isoformat(),
|
||
"runName": run_dir.name,
|
||
"outputDir": str(run_dir),
|
||
"baseDir": str(base_dir),
|
||
"objectName": model_filename,
|
||
"settings": settings,
|
||
"objectVertexCount": int(object_vertex_count),
|
||
"stats": stats,
|
||
"written": written,
|
||
}
|
||
|
||
|
||
def write_run_manifest(run_dir: Path, manifest: dict[str, Any]) -> Path:
|
||
run_dir.mkdir(parents=True, exist_ok=True)
|
||
path = run_dir / DATASET_RUN_FILENAME
|
||
path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
return path
|
||
|
||
|
||
def read_run_manifest(run_dir: Path) -> dict[str, Any] | None:
|
||
path = run_dir / DATASET_RUN_FILENAME
|
||
if not path.is_file():
|
||
return None
|
||
try:
|
||
data = json.loads(path.read_text(encoding="utf-8"))
|
||
except (OSError, json.JSONDecodeError):
|
||
return None
|
||
return data if isinstance(data, dict) else None
|
||
|
||
|
||
def _scene_entry_from_files(run_dir: Path, stem: str) -> dict[str, Any]:
|
||
npy_path = run_dir / f"{stem}.npy"
|
||
obj_path = run_dir / f"{stem}.obj"
|
||
entry: dict[str, Any] = {
|
||
"stem": stem,
|
||
"visibility": "unknown",
|
||
"hasObject": None,
|
||
"pointCount": None,
|
||
"objectPointCount": None,
|
||
}
|
||
if npy_path.is_file():
|
||
try:
|
||
rows = load_npy_float64_rows(npy_path)
|
||
object_point_count = sum(1 for row in rows if int(row[6]) == 1)
|
||
entry["pointCount"] = len(rows)
|
||
entry["objectPointCount"] = object_point_count
|
||
entry["hasObject"] = object_point_count > 0
|
||
except (OSError, ValueError, IndexError):
|
||
pass
|
||
elif obj_path.is_file():
|
||
try:
|
||
text = obj_path.read_text(encoding="utf-8", errors="ignore")
|
||
labeled = parse_obj_labeled_points(text)
|
||
if labeled:
|
||
object_point_count = sum(1 for row in labeled if int(row[3]) == 1)
|
||
entry["pointCount"] = len(labeled)
|
||
entry["objectPointCount"] = object_point_count
|
||
entry["hasObject"] = object_point_count > 0
|
||
else:
|
||
points = parse_obj_points(text)
|
||
entry["pointCount"] = len(points)
|
||
except OSError:
|
||
pass
|
||
return entry
|
||
|
||
|
||
def scan_run_scenes(run_dir: Path) -> list[dict[str, Any]]:
|
||
stems: set[str] = set()
|
||
for pattern in ("*.obj", "*.npy"):
|
||
for path in run_dir.glob(pattern):
|
||
if path.is_file():
|
||
stems.add(path.stem)
|
||
return [_scene_entry_from_files(run_dir, stem) for stem in sorted(stems)]
|
||
|
||
|
||
def _run_has_scene_files(path: Path) -> bool:
|
||
return any(path.glob("*.obj")) or any(path.glob("*.npy"))
|
||
|
||
|
||
def _relative_to_project(path: Path) -> str:
|
||
project_root = Path(__file__).resolve().parent.parent
|
||
try:
|
||
return str(path.relative_to(project_root))
|
||
except ValueError:
|
||
return str(path)
|
||
|
||
|
||
def list_dataset_runs(output_dir: str | Path = "sonar_dataset") -> list[dict[str, Any]]:
|
||
"""List dataset run folders under the base output directory."""
|
||
base = resolve_output_dir(output_dir)
|
||
runs: list[dict[str, Any]] = []
|
||
|
||
if base.is_dir() and _run_has_scene_files(base):
|
||
manifest = read_run_manifest(base)
|
||
written = manifest.get("written") if manifest else scan_run_scenes(base)
|
||
runs.append(
|
||
{
|
||
"runName": "(корень)",
|
||
"outputDir": str(base),
|
||
"loadPath": _relative_to_project(base),
|
||
"sceneCount": len(written),
|
||
"hasSettings": manifest is not None,
|
||
"generatedAt": manifest.get("generatedAt") if manifest else None,
|
||
"objectName": (manifest or {}).get("objectName")
|
||
or (manifest or {}).get("settings", {}).get("objectName"),
|
||
}
|
||
)
|
||
|
||
if not base.is_dir():
|
||
return runs
|
||
|
||
for child in sorted(base.iterdir(), key=lambda p: p.name, reverse=True):
|
||
if not child.is_dir() or not _run_has_scene_files(child):
|
||
continue
|
||
manifest = read_run_manifest(child)
|
||
written = manifest.get("written") if manifest else scan_run_scenes(child)
|
||
runs.append(
|
||
{
|
||
"runName": child.name,
|
||
"outputDir": str(child),
|
||
"loadPath": _relative_to_project(child),
|
||
"sceneCount": len(written),
|
||
"hasSettings": manifest is not None,
|
||
"generatedAt": manifest.get("generatedAt") if manifest else None,
|
||
"objectName": (manifest or {}).get("objectName")
|
||
or (manifest or {}).get("settings", {}).get("objectName"),
|
||
}
|
||
)
|
||
return runs
|
||
|
||
|
||
def load_dataset_run(output_dir: str | Path) -> dict[str, Any]:
|
||
"""Load a dataset run folder: settings, stats and scene list."""
|
||
run_dir = resolve_output_dir(output_dir)
|
||
if not run_dir.is_dir():
|
||
raise FileNotFoundError(f"Dataset folder not found: {run_dir}")
|
||
if not _run_has_scene_files(run_dir):
|
||
raise FileNotFoundError(f"No scene files in dataset folder: {run_dir}")
|
||
|
||
manifest = read_run_manifest(run_dir)
|
||
written = manifest.get("written") if manifest else scan_run_scenes(run_dir)
|
||
if not written:
|
||
raise FileNotFoundError(f"No scenes found in dataset folder: {run_dir}")
|
||
|
||
settings = dict((manifest or {}).get("settings") or {})
|
||
model_filename = (
|
||
settings.get("objectName")
|
||
or (manifest.get("objectName") if manifest else None)
|
||
)
|
||
stats = dict((manifest or {}).get("stats") or {})
|
||
if not stats:
|
||
stats = {
|
||
"total": len(written),
|
||
"withObject": sum(1 for item in written if item.get("hasObject")),
|
||
"withoutObject": sum(1 for item in written if item.get("hasObject") is False),
|
||
}
|
||
|
||
base_dir = Path(manifest["baseDir"]) if manifest and manifest.get("baseDir") else run_dir.parent
|
||
return {
|
||
"outputDir": str(run_dir),
|
||
"baseDir": str(base_dir),
|
||
"runName": manifest.get("runName") if manifest else run_dir.name,
|
||
"generatedAt": manifest.get("generatedAt") if manifest else None,
|
||
"hasSettings": manifest is not None,
|
||
"settings": settings,
|
||
"objectVertexCount": (manifest or {}).get("objectVertexCount"),
|
||
"stats": stats,
|
||
"written": written,
|
||
"count": settings.get("count") or len(written),
|
||
"seed": settings.get("seed"),
|
||
"beamCount": settings.get("beamCount"),
|
||
"lengthCount": settings.get("lengthCount"),
|
||
"objectScale": settings.get("objectScale"),
|
||
"objectScaleIsMax": settings.get("objectScaleIsMax"),
|
||
"objectName": model_filename,
|
||
"classLabels": {"0": "background", "1": "object"},
|
||
}
|
||
|
||
|
||
def _downsample_points(points: list[list[float]], max_points: int) -> list[list[float]]:
|
||
max_points = max(100, int(max_points))
|
||
if len(points) <= max_points:
|
||
return points
|
||
step = max(1, len(points) // max_points)
|
||
return points[::step][:max_points]
|
||
|
||
|
||
def load_npy_float64_rows(path: Path) -> list[list[float]]:
|
||
"""Read float64 little-endian .npy array written by export_npy_float64."""
|
||
import re
|
||
import struct
|
||
|
||
data = path.read_bytes()
|
||
if data[:6] != b"\x93NUMPY":
|
||
raise ValueError(f"Not a NumPy .npy file: {path.name}")
|
||
major = data[6]
|
||
if major == 1:
|
||
hlen = struct.unpack_from("<H", data, 8)[0]
|
||
header = data[10 : 10 + hlen].decode("latin1")
|
||
offset = 10 + hlen
|
||
elif major == 2:
|
||
hlen = struct.unpack_from("<I", data, 8)[0]
|
||
header = data[12 : 12 + hlen].decode("latin1")
|
||
offset = 12 + hlen
|
||
else:
|
||
raise ValueError(f"Unsupported .npy version: {major}")
|
||
match = re.search(r"shape'\s*:\s*\((\d+)\s*,\s*(\d+)\)", header)
|
||
if not match:
|
||
match = re.search(r"shape':\s*\((\d+),\s*(\d+)\)", header)
|
||
if not match:
|
||
raise ValueError(f"Cannot parse .npy shape from {path.name}")
|
||
n_rows, n_cols = int(match.group(1)), int(match.group(2))
|
||
expected = n_rows * n_cols * 8
|
||
body = data[offset : offset + expected]
|
||
if len(body) < expected:
|
||
raise ValueError(f"Truncated .npy payload in {path.name}")
|
||
flat = struct.unpack("<" + "d" * (n_rows * n_cols), body)
|
||
return [list(flat[i * n_cols : (i + 1) * n_cols]) for i in range(n_rows)]
|
||
|
||
|
||
def _class_counts_from_labeled(points: list[list[float]]) -> dict[str, int]:
|
||
"""Count classes from rows shaped [x, y, z, class]."""
|
||
counts: dict[str, int] = {}
|
||
for row in points:
|
||
key = str(int(round(float(row[3] if len(row) > 3 else 0))))
|
||
counts[key] = counts.get(key, 0) + 1
|
||
return counts
|
||
|
||
|
||
def _class_counts(rows: list[list[float]]) -> dict[str, int]:
|
||
counts: dict[str, int] = {}
|
||
for row in rows:
|
||
key = str(int(round(float(row[6] if len(row) > 6 else 0))))
|
||
counts[key] = counts.get(key, 0) + 1
|
||
return counts
|
||
|
||
|
||
def load_scene_preview(
|
||
*,
|
||
stem: str,
|
||
output_dir: str | Path = "sonar_dataset",
|
||
max_points: int = 25000,
|
||
) -> dict[str, Any]:
|
||
"""Load labeled points for a written scene (prefer .npy) for the 3D viewer.
|
||
|
||
Each preview point is [x, y, z, class].
|
||
"""
|
||
safe = "".join(ch if ch.isalnum() or ch in "_-" else "" for ch in (stem or ""))
|
||
if not safe or safe != stem:
|
||
raise ValueError("Invalid scene stem.")
|
||
out = resolve_output_dir(output_dir)
|
||
npy_path = out / f"{safe}.npy"
|
||
obj_path = out / f"{safe}.obj"
|
||
|
||
labeled: list[list[float]]
|
||
if npy_path.is_file():
|
||
rows = load_npy_float64_rows(npy_path)
|
||
labeled = [[float(r[0]), float(r[1]), float(r[2]), float(r[6])] for r in rows]
|
||
elif obj_path.is_file():
|
||
text = obj_path.read_text(encoding="utf-8", errors="ignore")
|
||
labeled = parse_obj_labeled_points(text)
|
||
if not labeled:
|
||
points = parse_obj_points(text)
|
||
labeled = [[p[0], p[1], p[2], 0.0] for p in points]
|
||
else:
|
||
raise FileNotFoundError(f"Scene not found: {safe}.npy / {safe}.obj")
|
||
|
||
full_counts = _class_counts_from_labeled(labeled)
|
||
preview = _downsample_points(labeled, max_points)
|
||
return {
|
||
"stem": safe,
|
||
"outputDir": str(out),
|
||
"pointCount": len(labeled),
|
||
"previewCount": len(preview),
|
||
"points": preview,
|
||
"classCounts": full_counts,
|
||
"classLabels": {"0": "background", "1": "object"},
|
||
"obj": str(obj_path) if obj_path.is_file() else None,
|
||
"npy": str(npy_path) if npy_path.is_file() else None,
|
||
}
|
||
|
||
|
||
def write_scene_files(
|
||
scene: dict[str, Any],
|
||
output_dir: Path,
|
||
stem: str,
|
||
) -> dict[str, str]:
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
npy_path = output_dir / f"{stem}.npy"
|
||
obj_path = output_dir / f"{stem}.obj"
|
||
npy_path.write_bytes(export_npy_float64(scene["rows"]))
|
||
classes = [int(r[6]) for r in scene["rows"]]
|
||
obj_path.write_text(
|
||
export_obj(scene["points"], object_name=stem, classes=classes),
|
||
encoding="utf-8",
|
||
)
|
||
return {"npy": str(npy_path), "obj": str(obj_path), "stem": stem}
|
||
|
||
|
||
def iter_generate_dataset(
|
||
*,
|
||
count: int = 5,
|
||
seed: int = 42,
|
||
output_dir: str | Path = "sonar_dataset",
|
||
object_points: list[list[float]] | None = None,
|
||
object_name: str | None = None,
|
||
object_kind: str | None = None,
|
||
object_scale: float = 1.0,
|
||
object_scale_is_max: bool = False,
|
||
beam_count: int = 45,
|
||
length_count: int | None = None,
|
||
absent_pct: float = 30.0,
|
||
nearly_hidden_pct: float = 20.0,
|
||
partial_pct: float = 40.0,
|
||
visible_pct: float = 40.0,
|
||
mle_mode: bool = False,
|
||
auv_depth_min: float = 2.0,
|
||
auv_depth_max: float = 8.0,
|
||
relief_scale_pct: float = 100.0,
|
||
):
|
||
"""Yield NDJSON-friendly progress events, then a final ``done`` payload.
|
||
|
||
Events:
|
||
{"type":"start","total":N,"outputDir":"...","runName":"..."}
|
||
{"type":"progress","current":k,"total":N,"entry":{...}}
|
||
{"type":"done","result":{...}}
|
||
"""
|
||
count = int(count)
|
||
if count < 1:
|
||
raise ValueError("count must be >= 1")
|
||
if count > 5000:
|
||
raise ValueError("count must be <= 5000")
|
||
kind = (object_kind or "").strip().lower() or None
|
||
if kind == "pipe":
|
||
template: list[list[float]] = []
|
||
else:
|
||
if not object_points or len(object_points) < 3:
|
||
raise ValueError("A valid .obj model with at least 3 vertices is required.")
|
||
template = normalize_object_points(object_points)
|
||
object_scale = float(object_scale)
|
||
if object_scale <= 0:
|
||
raise ValueError("object_scale must be > 0")
|
||
if object_scale > 100:
|
||
raise ValueError("object_scale must be <= 100")
|
||
object_scale_is_max = bool(object_scale_is_max)
|
||
beam_count = int(beam_count)
|
||
if beam_count < 1:
|
||
raise ValueError("beam_count (Кол-во лучей) must be >= 1")
|
||
if beam_count > 1024:
|
||
raise ValueError("beam_count (Кол-во лучей) must be <= 1024")
|
||
if length_count is None:
|
||
length_count = beam_count
|
||
length_count = int(length_count)
|
||
if length_count < 1:
|
||
raise ValueError("length_count (Длина) must be >= 1")
|
||
if length_count > 1024:
|
||
raise ValueError("length_count (Длина) must be <= 1024")
|
||
|
||
use_mle = bool(mle_mode)
|
||
depth_lo = max(0.3, float(auv_depth_min))
|
||
depth_hi = max(0.3, float(auv_depth_max))
|
||
if depth_hi < depth_lo:
|
||
depth_lo, depth_hi = depth_hi, depth_lo
|
||
if use_mle and (not math.isfinite(depth_lo) or not math.isfinite(depth_hi)):
|
||
raise ValueError("Диапазон высот must be finite numbers")
|
||
|
||
relief_pct = max(5.0, min(500.0, float(relief_scale_pct)))
|
||
if not math.isfinite(relief_pct):
|
||
raise ValueError("relief_scale_pct (Размер неровностей) must be finite")
|
||
|
||
absent_pct = _clamp_pct(absent_pct, name="absent_pct")
|
||
nearly_hidden_pct = _clamp_pct(nearly_hidden_pct, name="nearly_hidden_pct")
|
||
partial_pct = _clamp_pct(partial_pct, name="partial_pct")
|
||
visible_pct = _clamp_pct(visible_pct, name="visible_pct")
|
||
|
||
object_name = _normalize_model_filename(object_name)
|
||
|
||
base = resolve_output_dir(output_dir)
|
||
run_dir = make_generation_run_dir(base, object_name)
|
||
|
||
labels = plan_scene_labels(
|
||
count,
|
||
seed,
|
||
absent_pct=absent_pct,
|
||
nearly_hidden_pct=nearly_hidden_pct,
|
||
partial_pct=partial_pct,
|
||
visible_pct=visible_pct,
|
||
)
|
||
written: list[dict[str, Any]] = []
|
||
stats = {
|
||
"total": count,
|
||
"withObject": 0,
|
||
"withoutObject": 0,
|
||
"nearly_hidden": 0,
|
||
"partial": 0,
|
||
"visible": 0,
|
||
"absent": 0,
|
||
}
|
||
|
||
preview_points: list[list[float]] | None = None
|
||
preview_stem: str | None = None
|
||
preview_has_object = False
|
||
object_vertex_count = len(template)
|
||
|
||
yield {
|
||
"type": "start",
|
||
"total": count,
|
||
"outputDir": str(run_dir),
|
||
"baseDir": str(base),
|
||
"runName": run_dir.name,
|
||
}
|
||
|
||
for i in range(count):
|
||
visibility = labels[i]
|
||
scene_seed = int(seed) + i * 10007 + 17
|
||
scene_rng = random.Random(scene_seed ^ 0xC0FFEE)
|
||
if object_scale_is_max:
|
||
lo, hi = 1.0, object_scale
|
||
if hi < lo:
|
||
lo, hi = hi, lo
|
||
scene_scale = scene_rng.uniform(lo, hi)
|
||
else:
|
||
scene_scale = object_scale
|
||
scene_depth = scene_rng.uniform(depth_lo, depth_hi) if use_mle else None
|
||
scene = generate_sonar_scene(
|
||
seed=scene_seed,
|
||
visibility=visibility,
|
||
object_points=template,
|
||
object_scale=scene_scale,
|
||
object_kind=kind,
|
||
beam_count=beam_count,
|
||
length_count=length_count,
|
||
mle_mode=use_mle,
|
||
auv_depth=scene_depth,
|
||
relief_scale_pct=relief_pct,
|
||
)
|
||
if scene.get("object") and scene["object"].get("vertexCount"):
|
||
object_vertex_count = int(scene["object"]["vertexCount"])
|
||
area, scene_no, stem = scene_index_to_area_name(i)
|
||
paths = write_scene_files(scene, run_dir, stem)
|
||
|
||
entry = {
|
||
"index": i,
|
||
"area": area,
|
||
"scene": scene_no,
|
||
"stem": stem,
|
||
"visibility": visibility,
|
||
"hasObject": scene["hasObject"],
|
||
"pointCount": scene["pointCount"],
|
||
"objectPointCount": scene["objectPointCount"],
|
||
"objectScale": scene_scale,
|
||
"auvDepth": scene.get("auvDepth"),
|
||
"files": paths,
|
||
}
|
||
written.append(entry)
|
||
|
||
stats[visibility] = stats.get(visibility, 0) + 1
|
||
if scene["hasObject"]:
|
||
stats["withObject"] += 1
|
||
else:
|
||
stats["withoutObject"] += 1
|
||
|
||
if preview_points is None or (scene["hasObject"] and not preview_has_object):
|
||
preview_points = [[r[0], r[1], r[2], r[6]] for r in scene["rows"]]
|
||
preview_stem = stem
|
||
preview_has_object = bool(scene["hasObject"])
|
||
|
||
yield {
|
||
"type": "progress",
|
||
"current": i + 1,
|
||
"total": count,
|
||
"entry": entry,
|
||
"outputDir": str(run_dir),
|
||
}
|
||
|
||
preview: dict[str, Any] | None = None
|
||
if preview_points is not None:
|
||
pts = _downsample_points(preview_points, 25000)
|
||
preview = {
|
||
"stem": preview_stem,
|
||
"points": pts,
|
||
"pointCount": len(preview_points),
|
||
"classCounts": _class_counts_from_labeled(preview_points),
|
||
"classLabels": {"0": "background", "1": "object"},
|
||
}
|
||
|
||
result = {
|
||
"outputDir": str(run_dir),
|
||
"baseDir": str(base),
|
||
"runName": run_dir.name,
|
||
"count": count,
|
||
"seed": int(seed),
|
||
"beamCount": beam_count,
|
||
"lengthCount": length_count,
|
||
"mleMode": use_mle,
|
||
"auvDepthMin": depth_lo if use_mle else None,
|
||
"auvDepthMax": depth_hi if use_mle else None,
|
||
"reliefScalePct": relief_pct,
|
||
"objectName": object_name,
|
||
"objectKind": kind,
|
||
"objectScale": object_scale,
|
||
"objectScaleIsMax": object_scale_is_max,
|
||
"objectVertexCount": object_vertex_count,
|
||
"absentPct": absent_pct,
|
||
"nearlyHiddenPct": nearly_hidden_pct,
|
||
"partialPct": partial_pct,
|
||
"visiblePct": visible_pct,
|
||
"classLabels": {"0": "background", "1": "object"},
|
||
"stats": stats,
|
||
"written": written,
|
||
"preview": preview,
|
||
}
|
||
manifest = build_run_manifest(
|
||
run_dir=run_dir,
|
||
base_dir=base,
|
||
count=count,
|
||
seed=int(seed),
|
||
output_dir=str(output_dir),
|
||
object_name=object_name,
|
||
object_scale=object_scale,
|
||
object_scale_is_max=object_scale_is_max,
|
||
beam_count=beam_count,
|
||
length_count=length_count,
|
||
object_vertex_count=len(template),
|
||
stats=stats,
|
||
written=written,
|
||
absent_pct=absent_pct,
|
||
nearly_hidden_pct=nearly_hidden_pct,
|
||
partial_pct=partial_pct,
|
||
visible_pct=visible_pct,
|
||
mle_mode=use_mle,
|
||
auv_depth_min=depth_lo if use_mle else None,
|
||
auv_depth_max=depth_hi if use_mle else None,
|
||
relief_scale_pct=relief_pct,
|
||
)
|
||
write_run_manifest(run_dir, manifest)
|
||
result["settingsPath"] = str(run_dir / DATASET_RUN_FILENAME)
|
||
yield {"type": "done", "result": result}
|
||
|
||
|
||
def generate_dataset(
|
||
*,
|
||
count: int = 5,
|
||
seed: int = 42,
|
||
output_dir: str | Path = "sonar_dataset",
|
||
object_points: list[list[float]] | None = None,
|
||
object_name: str | None = None,
|
||
object_kind: str | None = None,
|
||
object_scale: float = 1.0,
|
||
object_scale_is_max: bool = False,
|
||
beam_count: int = 45,
|
||
length_count: int | None = None,
|
||
absent_pct: float = 30.0,
|
||
nearly_hidden_pct: float = 20.0,
|
||
partial_pct: float = 40.0,
|
||
visible_pct: float = 40.0,
|
||
mle_mode: bool = False,
|
||
auv_depth_min: float = 2.0,
|
||
auv_depth_max: float = 8.0,
|
||
relief_scale_pct: float = 100.0,
|
||
) -> dict[str, Any]:
|
||
"""Generate `count` unique scenes into a new timestamped run folder under output_dir."""
|
||
result: dict[str, Any] | None = None
|
||
for event in iter_generate_dataset(
|
||
count=count,
|
||
seed=seed,
|
||
output_dir=output_dir,
|
||
object_points=object_points,
|
||
object_name=object_name,
|
||
object_kind=object_kind,
|
||
object_scale=object_scale,
|
||
object_scale_is_max=object_scale_is_max,
|
||
beam_count=beam_count,
|
||
length_count=length_count,
|
||
absent_pct=absent_pct,
|
||
nearly_hidden_pct=nearly_hidden_pct,
|
||
partial_pct=partial_pct,
|
||
visible_pct=visible_pct,
|
||
mle_mode=mle_mode,
|
||
auv_depth_min=auv_depth_min,
|
||
auv_depth_max=auv_depth_max,
|
||
relief_scale_pct=relief_scale_pct,
|
||
):
|
||
if event.get("type") == "done":
|
||
result = event["result"]
|
||
if result is None:
|
||
raise RuntimeError("Dataset generation produced no result.")
|
||
return result
|