Реструктуризация проекта и генератор синтетических датасетов эхолота.
Перенесены backend/frontend/desktop/engine, добавлены вкладки конструктора сцен и генератора датасета с параметрами лучей и длины сетки рельефа, обновлены API и Docker-сборка. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Cursor
@@ -0,0 +1,785 @@
|
||||
"""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,
|
||||
}
|
||||
Reference in New Issue
Block a user