gitrusprusandCursor
6cfc16e53c
Вкладки позволяют готовить рельеф, двигать АНПА, накапливать поверхность по лучам и сохранять скриншоты окон; ГБО использует бортовые секторы 12–75° и чёрные зоны вне обзора. Co-authored-by: Cursor <cursoragent@cursor.com>
1056 lines
36 KiB
Python
1056 lines
36 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,
|
||
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")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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
|
||
# ---------------------------------------------------------------------------
|
||
|
||
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,
|
||
) -> dict[str, Any]:
|
||
"""Build continuous seafloor heightfield parameters (no point cloud yet)."""
|
||
beams = max(1, int(beam_count))
|
||
length_points = beams if length_count is None else max(1, int(length_count))
|
||
size_x = rng.uniform(8.0, 16.0)
|
||
size_y = rng.uniform(8.0, 16.0)
|
||
base_z = rng.uniform(-1.2, -0.2)
|
||
amplitude = rng.uniform(0.05, 0.35)
|
||
frequency = rng.uniform(0.4, 2.2)
|
||
|
||
hills = [
|
||
(
|
||
rng.uniform(-size_x * 0.4, size_x * 0.4),
|
||
rng.uniform(-size_y * 0.4, size_y * 0.4),
|
||
rng.uniform(0.15, 0.7),
|
||
rng.uniform(0.6, 2.2),
|
||
)
|
||
for _ in range(rng.randint(1, 4))
|
||
]
|
||
valleys = [
|
||
(
|
||
rng.uniform(-size_x * 0.4, size_x * 0.4),
|
||
rng.uniform(-size_y * 0.4, size_y * 0.4),
|
||
rng.uniform(0.1, 0.55),
|
||
rng.uniform(0.5, 2.0),
|
||
)
|
||
for _ in range(rng.randint(1, 3))
|
||
]
|
||
# Relief clutter / false features as heightfield bumps (not extra points)
|
||
bumps = [
|
||
(
|
||
rng.uniform(-size_x * 0.45, size_x * 0.45),
|
||
rng.uniform(-size_y * 0.45, size_y * 0.45),
|
||
rng.uniform(0.03, 0.35),
|
||
rng.uniform(0.15, 0.9),
|
||
)
|
||
for _ in range(rng.randint(4, 14))
|
||
]
|
||
|
||
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),
|
||
}
|
||
|
||
|
||
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 _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
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Balance plan + single scene
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def plan_scene_labels(count: int, seed: int) -> list[str]:
|
||
"""Return visibility label per scene: absent | nearly_hidden | partial | visible.
|
||
|
||
~50% absent; among present scenes, roughly equal nearly_hidden/partial/visible.
|
||
"""
|
||
count = max(0, int(count))
|
||
rng = random.Random(int(seed) ^ 0xA5A5_5A5A)
|
||
n_with = (count + 1) // 2 # ceil → ~50% with object
|
||
n_without = count - n_with
|
||
labels: list[str] = ["absent"] * n_without
|
||
for i in range(n_with):
|
||
labels.append(VISIBILITY_TIERS[i % 3])
|
||
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,
|
||
beam_count: int = 45,
|
||
length_count: int | None = None,
|
||
) -> 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.
|
||
"""
|
||
rng = random.Random(int(seed))
|
||
if visibility not in ("absent",) + VISIBILITY_TIERS:
|
||
raise ValueError(f"Unknown visibility: {visibility}")
|
||
if visibility != "absent" and not object_points:
|
||
raise ValueError("object_points required when visibility is not absent.")
|
||
|
||
meta = _build_seafloor_meta(rng, beam_count=beam_count, length_count=length_count)
|
||
expected = int(meta["beamCount"]) * int(meta["lengthCount"])
|
||
|
||
object_info: dict[str, Any] | None = None
|
||
object_hits: dict[tuple[int, int], float] | None = None
|
||
if visibility != "absent":
|
||
world, object_info = _place_object_in_scene(
|
||
rng,
|
||
meta,
|
||
visibility,
|
||
object_points,
|
||
object_scale=object_scale,
|
||
)
|
||
object_hits = _rasterize_object_hits(world, meta)
|
||
object_info["hitCellCount"] = len(object_hits)
|
||
object_info["keptCount"] = len(object_hits)
|
||
object_info["requestedCount"] = expected
|
||
|
||
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,
|
||
"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,
|
||
"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]],
|
||
) -> dict[str, Any]:
|
||
model_filename = _normalize_model_filename(object_name)
|
||
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": {
|
||
"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,
|
||
},
|
||
"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]],
|
||
object_name: str | None = None,
|
||
object_scale: float = 1.0,
|
||
object_scale_is_max: bool = False,
|
||
beam_count: int = 45,
|
||
length_count: int | None = None,
|
||
):
|
||
"""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")
|
||
if not object_points or len(object_points) < 3:
|
||
raise ValueError("A valid .obj model with at least 3 vertices is required.")
|
||
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")
|
||
|
||
object_name = _normalize_model_filename(object_name)
|
||
|
||
base = resolve_output_dir(output_dir)
|
||
run_dir = make_generation_run_dir(base, object_name)
|
||
template = normalize_object_points(object_points)
|
||
|
||
labels = plan_scene_labels(count, seed)
|
||
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
|
||
|
||
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 = generate_sonar_scene(
|
||
seed=scene_seed,
|
||
visibility=visibility,
|
||
object_points=template,
|
||
object_scale=scene_scale,
|
||
beam_count=beam_count,
|
||
length_count=length_count,
|
||
)
|
||
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,
|
||
"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,
|
||
"objectName": object_name,
|
||
"objectScale": object_scale,
|
||
"objectScaleIsMax": object_scale_is_max,
|
||
"objectVertexCount": len(template),
|
||
"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,
|
||
)
|
||
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]],
|
||
object_name: str | None = None,
|
||
object_scale: float = 1.0,
|
||
object_scale_is_max: bool = False,
|
||
beam_count: int = 45,
|
||
length_count: int | None = None,
|
||
) -> 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_scale=object_scale,
|
||
object_scale_is_max=object_scale_is_max,
|
||
beam_count=beam_count,
|
||
length_count=length_count,
|
||
):
|
||
if event.get("type") == "done":
|
||
result = event["result"]
|
||
if result is None:
|
||
raise RuntimeError("Dataset generation produced no result.")
|
||
return result
|