gitrusprusandCursor
4f253b860f
Перенесены backend/frontend/desktop/engine, добавлены вкладки конструктора сцен и генератора датасета с параметрами лучей и длины сетки рельефа, обновлены API и Docker-сборка. Co-authored-by: Cursor <cursoragent@cursor.com>
786 lines
26 KiB
Python
786 lines
26 KiB
Python
"""Batch synthetic sonar dataset generator for PointNet semantic segmentation.
|
|
|
|
Produces paired Area_X_scene_XXXX.npy + .obj files under sonar_dataset/.
|
|
Target class 1 = user-provided object (from .obj mesh vertices); class 0 = seafloor / clutter.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import random
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from scene_generator import (
|
|
apply_transform,
|
|
export_npy_float64,
|
|
export_obj,
|
|
generate_box,
|
|
generate_pipe,
|
|
generate_sphere,
|
|
generate_torus,
|
|
parse_obj_points,
|
|
points_to_pointnet_rows,
|
|
)
|
|
|
|
# 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 resample_object_points(
|
|
template: list[list[float]],
|
|
count: int,
|
|
*,
|
|
noise: float = 0.0,
|
|
seed: int = 1,
|
|
) -> list[list[float]]:
|
|
"""Subsample (or sample with replacement) template points to the requested count."""
|
|
if not template:
|
|
raise ValueError("Object template is empty.")
|
|
rng = random.Random(int(seed))
|
|
count = max(1, int(count))
|
|
out: list[list[float]] = []
|
|
n = len(template)
|
|
for _ in range(count):
|
|
src = template[rng.randrange(n)]
|
|
if noise > 0:
|
|
out.append(
|
|
[
|
|
src[0] + rng.uniform(-noise, noise),
|
|
src[1] + rng.uniform(-noise, noise),
|
|
src[2] + rng.uniform(-noise, noise),
|
|
]
|
|
)
|
|
else:
|
|
out.append([src[0], src[1], src[2]])
|
|
return out
|
|
|
|
|
|
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 / clutter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
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 _generate_seafloor(
|
|
rng: random.Random,
|
|
*,
|
|
beam_count: int = 45,
|
|
length_count: int | None = None,
|
|
) -> tuple[list[list[float]], dict[str, Any]]:
|
|
"""Sample seafloor as a square relief grid.
|
|
|
|
beam_count controls width resolution (X axis).
|
|
length_count controls length resolution (Y axis).
|
|
"""
|
|
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)
|
|
noise = rng.uniform(0.005, 0.04)
|
|
|
|
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))
|
|
]
|
|
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.18),
|
|
rng.uniform(0.15, 0.55),
|
|
)
|
|
for _ in range(rng.randint(3, 12))
|
|
]
|
|
|
|
meta = {
|
|
"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,
|
|
}
|
|
|
|
half_x = size_x * 0.5
|
|
half_y = size_y * 0.5
|
|
points: list[list[float]] = []
|
|
for yi in range(length_points):
|
|
y = -half_y if length_points == 1 else (-half_y + size_y * yi / (length_points - 1))
|
|
for xi in range(beams):
|
|
x = -half_x if beams == 1 else (-half_x + size_x * xi / (beams - 1))
|
|
x += rng.uniform(-noise * 2, noise * 2)
|
|
yj = y + rng.uniform(-noise * 2, noise * 2)
|
|
z = _seafloor_height(
|
|
x,
|
|
yj,
|
|
base_z=base_z,
|
|
amplitude=amplitude,
|
|
frequency=frequency,
|
|
hills=hills,
|
|
valleys=valleys,
|
|
bumps=bumps,
|
|
)
|
|
z += rng.uniform(-noise, noise)
|
|
points.append([x, yj, z])
|
|
|
|
# Local noise clusters (false sonar clutter blobs)
|
|
for _ in range(rng.randint(1, 5)):
|
|
cx = rng.uniform(-half_x * 0.8, half_x * 0.8)
|
|
cy = rng.uniform(-half_y * 0.8, half_y * 0.8)
|
|
cz = _seafloor_height(
|
|
cx,
|
|
cy,
|
|
base_z=base_z,
|
|
amplitude=amplitude,
|
|
frequency=frequency,
|
|
hills=hills,
|
|
valleys=valleys,
|
|
bumps=bumps,
|
|
) + rng.uniform(0.0, 0.25)
|
|
n_blob = rng.randint(40, 280)
|
|
spread = rng.uniform(0.15, 0.7)
|
|
for _ in range(n_blob):
|
|
points.append(
|
|
[
|
|
cx + rng.gauss(0, spread),
|
|
cy + rng.gauss(0, spread),
|
|
cz + rng.gauss(0, spread * 0.35),
|
|
]
|
|
)
|
|
|
|
meta["pingCount"] = length_points
|
|
meta["swathBeams"] = beams
|
|
return points, meta
|
|
|
|
|
|
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 _generate_false_objects(rng: random.Random, meta: dict[str, Any]) -> list[list[float]]:
|
|
n_objects = rng.randint(0, 6)
|
|
points: list[list[float]] = []
|
|
half_x = float(meta["sizeX"]) * 0.5
|
|
half_y = float(meta["sizeY"]) * 0.5
|
|
|
|
for i in range(n_objects):
|
|
kind = rng.choice(["sphere", "box", "torus", "pipe"])
|
|
count = rng.randint(80, 900)
|
|
noise = rng.uniform(0.005, 0.03)
|
|
seed = rng.randint(0, 10_000_000)
|
|
if kind == "sphere":
|
|
local = generate_sphere(
|
|
{"radius": rng.uniform(0.08, 0.55), "count": count, "noise": noise, "seed": seed}
|
|
)
|
|
elif kind == "box":
|
|
local = generate_box(
|
|
{
|
|
"sizeX": rng.uniform(0.15, 1.2),
|
|
"sizeY": rng.uniform(0.15, 1.0),
|
|
"sizeZ": rng.uniform(0.08, 0.6),
|
|
"count": count,
|
|
"noise": noise,
|
|
"seed": seed,
|
|
}
|
|
)
|
|
elif kind == "torus":
|
|
major = rng.uniform(0.15, 0.6)
|
|
local = generate_torus(
|
|
{
|
|
"majorR": major,
|
|
"minorR": rng.uniform(0.03, major * 0.4),
|
|
"count": count,
|
|
"noise": noise,
|
|
"seed": seed,
|
|
}
|
|
)
|
|
else:
|
|
local = generate_pipe(
|
|
{
|
|
"length": rng.uniform(0.4, 2.5),
|
|
"radius": rng.uniform(0.04, 0.2),
|
|
"axis": rng.choice(["x", "y", "z"]),
|
|
"count": count,
|
|
"noise": noise,
|
|
"seed": seed,
|
|
}
|
|
)
|
|
|
|
tx = rng.uniform(-half_x * 0.75, half_x * 0.75)
|
|
ty = rng.uniform(-half_y * 0.75, half_y * 0.75)
|
|
floor_z = _height_at(tx, ty, meta)
|
|
# Rest on / slightly into seafloor
|
|
tz = floor_z + rng.uniform(-0.05, 0.35)
|
|
transform = {
|
|
"x": tx,
|
|
"y": ty,
|
|
"z": tz,
|
|
"rx": rng.uniform(-0.4, 0.4),
|
|
"ry": rng.uniform(-0.4, 0.4),
|
|
"rz": rng.uniform(0, 2 * math.pi),
|
|
}
|
|
world = apply_transform(local, transform)
|
|
# Drop points buried deep under seafloor
|
|
for p in world:
|
|
if p[2] >= _height_at(p[0], p[1], meta) - 0.02:
|
|
points.append(p)
|
|
return points
|
|
|
|
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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 _place_object(
|
|
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]]:
|
|
"""Sample, transform, and bury target object; return surviving world points + info."""
|
|
base_scale = max(0.01, float(object_scale))
|
|
if visibility == "nearly_hidden":
|
|
count = rng.randint(80, 600)
|
|
burial = rng.uniform(0.35, 0.75)
|
|
scale = base_scale * rng.uniform(0.7, 1.15)
|
|
elif visibility == "partial":
|
|
count = rng.randint(400, 2500)
|
|
burial = rng.uniform(0.12, 0.4)
|
|
scale = base_scale * rng.uniform(0.8, 1.3)
|
|
else: # visible
|
|
count = rng.randint(1500, 8000)
|
|
burial = rng.uniform(-0.05, 0.15)
|
|
scale = base_scale * rng.uniform(0.85, 1.4)
|
|
|
|
noise = rng.uniform(0.004, 0.025)
|
|
local = resample_object_points(
|
|
object_template,
|
|
count,
|
|
noise=noise,
|
|
seed=rng.randint(0, 10_000_000),
|
|
)
|
|
# Apply world scale to unit-normalized template
|
|
local = [[p[0] * scale, p[1] * scale, p[2] * scale] for p in local]
|
|
|
|
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)
|
|
|
|
kept: list[list[float]] = []
|
|
for p in world:
|
|
surface = _height_at(p[0], p[1], meta)
|
|
eps = 0.01 if visibility != "nearly_hidden" else -0.02
|
|
if p[2] >= surface + eps:
|
|
kept.append(p)
|
|
|
|
if visibility == "nearly_hidden" and len(kept) < 15 and world:
|
|
ranked = sorted(world, key=lambda p: p[2] - _height_at(p[0], p[1], meta), reverse=True)
|
|
kept = ranked[: max(15, min(40, len(ranked) // 8))]
|
|
|
|
info = {
|
|
"visibility": visibility,
|
|
"transform": transform,
|
|
"requestedCount": count,
|
|
"keptCount": len(kept),
|
|
"scale": scale,
|
|
"objectScale": base_scale,
|
|
"burial": burial,
|
|
"classLabel": "object",
|
|
"classId": 1,
|
|
}
|
|
return kept, info
|
|
|
|
|
|
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 unique sonar scene. visibility in absent|nearly_hidden|partial|visible."""
|
|
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.")
|
|
|
|
floor_pts, meta = _generate_seafloor(rng, beam_count=beam_count, length_count=length_count)
|
|
clutter = _generate_false_objects(rng, meta)
|
|
|
|
jitter = rng.uniform(0.0, 0.015)
|
|
background = floor_pts + clutter
|
|
if jitter > 0:
|
|
background = [
|
|
[
|
|
p[0] + rng.uniform(-jitter, jitter),
|
|
p[1] + rng.uniform(-jitter, jitter),
|
|
p[2] + rng.uniform(-jitter, jitter),
|
|
]
|
|
for p in background
|
|
]
|
|
|
|
drop = rng.uniform(0.0, 0.12)
|
|
if drop > 0:
|
|
background = [p for p in background if rng.random() >= drop]
|
|
|
|
object_pts: list[list[float]] = []
|
|
object_info: dict[str, Any] | None = None
|
|
if visibility != "absent":
|
|
object_pts, object_info = _place_object(
|
|
rng,
|
|
meta,
|
|
visibility,
|
|
object_points,
|
|
object_scale=object_scale,
|
|
)
|
|
|
|
# class 0 = background, class 1 = object
|
|
rows = points_to_pointnet_rows(background, 0.0)
|
|
rows.extend(points_to_pointnet_rows(object_pts, 1.0))
|
|
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": len(object_pts),
|
|
"backgroundPointCount": len(background),
|
|
"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"),
|
|
"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
|
|
|
|
|
|
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")
|
|
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"]))
|
|
obj_path.write_text(export_obj(scene["points"], object_name=stem), encoding="utf-8")
|
|
return {"npy": str(npy_path), "obj": str(obj_path), "stem": stem}
|
|
|
|
|
|
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,
|
|
beam_count: int = 45,
|
|
length_count: int | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Generate `count` unique scenes into output_dir with Area_X naming.
|
|
|
|
object_points: normalized template vertices from user .obj (class 1 = object).
|
|
object_scale: relative size multiplier vs unit-normalized mesh (1.0 = default).
|
|
beam_count: number of width points for seafloor grid (X axis).
|
|
length_count: number of length points for seafloor grid (Y axis). Defaults to beam_count.
|
|
"""
|
|
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")
|
|
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")
|
|
|
|
out = resolve_output_dir(output_dir)
|
|
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
|
|
|
|
for i in range(count):
|
|
visibility = labels[i]
|
|
scene_seed = int(seed) + i * 10007 + 17
|
|
scene = generate_sonar_scene(
|
|
seed=scene_seed,
|
|
visibility=visibility,
|
|
object_points=template,
|
|
object_scale=object_scale,
|
|
beam_count=beam_count,
|
|
length_count=length_count,
|
|
)
|
|
area, scene_no, stem = scene_index_to_area_name(i)
|
|
paths = write_scene_files(scene, out, stem)
|
|
|
|
entry = {
|
|
"index": i,
|
|
"area": area,
|
|
"scene": scene_no,
|
|
"stem": stem,
|
|
"visibility": visibility,
|
|
"hasObject": scene["hasObject"],
|
|
"pointCount": scene["pointCount"],
|
|
"objectPointCount": scene["objectPointCount"],
|
|
"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"])
|
|
|
|
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"},
|
|
}
|
|
|
|
return {
|
|
"outputDir": str(out),
|
|
"count": count,
|
|
"seed": int(seed),
|
|
"beamCount": beam_count,
|
|
"lengthCount": length_count,
|
|
"objectName": object_name,
|
|
"objectScale": object_scale,
|
|
"objectVertexCount": len(template),
|
|
"classLabels": {"0": "background", "1": "object"},
|
|
"stats": stats,
|
|
"written": written,
|
|
"preview": preview,
|
|
}
|