gitrusprus/DotsToSirface
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

best version

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
gitrusprus
2026-07-17 15:06:51 +03:00
co-authored by Cursor
parent 4f253b860f
commit 643d834ba7
5 changed files with 310 additions and 289 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/dataset_generator.py
+219 -274
View File
@@ -1,7 +1,10 @@
"""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.
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
@@ -15,12 +18,7 @@ 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).
@@ -84,35 +82,6 @@ def load_object_points_from_obj_text(text: str) -> list[list[float]]:
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
@@ -121,7 +90,7 @@ def object_half_extent_z(points: list[list[float]]) -> float:
# ---------------------------------------------------------------------------
# Seafloor / clutter
# Seafloor heightfield + echosounder ray casting
# ---------------------------------------------------------------------------
def _seafloor_height(
@@ -153,17 +122,13 @@ def _seafloor_height(
return z
def _generate_seafloor(
def _build_seafloor_meta(
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).
"""
) -> 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)
@@ -171,7 +136,6 @@ def _generate_seafloor(
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 = [
(
@@ -191,17 +155,18 @@ def _generate_seafloor(
)
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.18),
rng.uniform(0.15, 0.55),
rng.uniform(0.03, 0.35),
rng.uniform(0.15, 0.9),
)
for _ in range(rng.randint(3, 12))
for _ in range(rng.randint(4, 14))
]
meta = {
return {
"sizeX": size_x,
"sizeY": size_y,
"baseZ": base_z,
@@ -214,59 +179,11 @@ def _generate_seafloor(
"lengthCount": length_points,
"gridWidthPoints": beams,
"gridLengthPoints": length_points,
"pingCount": length_points,
"swathBeams": beams,
"noise": rng.uniform(0.002, 0.02),
}
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(
@@ -281,76 +198,167 @@ def _height_at(x: float, y: float, meta: dict[str, Any]) -> float:
)
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
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
# ---------------------------------------------------------------------------
@@ -373,82 +381,6 @@ def plan_scene_labels(count: int, seed: int) -> list[str]:
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,
@@ -458,48 +390,42 @@ def generate_sonar_scene(
beam_count: int = 45,
length_count: int | None = None,
) -> dict[str, Any]:
"""Build one unique sonar scene. visibility in absent|nearly_hidden|partial|visible."""
"""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.")
floor_pts, meta = _generate_seafloor(rng, beam_count=beam_count, length_count=length_count)
clutter = _generate_false_objects(rng, meta)
meta = _build_seafloor_meta(rng, beam_count=beam_count, length_count=length_count)
expected = int(meta["beamCount"]) * int(meta["lengthCount"])
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
object_hits: dict[tuple[int, int], float] | None = None
if visibility != "absent":
object_pts, object_info = _place_object(
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)}")
# 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),
@@ -507,15 +433,18 @@ def generate_sonar_scene(
"hasObject": visibility != "absent",
"object": object_info,
"pointCount": len(rows),
"objectPointCount": len(object_pts),
"backgroundPointCount": len(background),
"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),
"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",
@@ -523,6 +452,7 @@ def generate_sonar_scene(
),
"pingCount": meta.get("pingCount"),
"swathBeams": meta.get("swathBeams"),
"gridPointCount": expected,
"floorFeatures": {
"hills": len(meta["hills"]),
"valleys": len(meta["valleys"]),
@@ -666,6 +596,7 @@ def generate_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]:
@@ -673,8 +604,11 @@ def generate_dataset(
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.
object_scale_is_max: if True, treat object_scale as upper bound and sample
per-scene scale uniformly from [1, object_scale] (AUV altitude variation).
beam_count: across-track beams (width resolution, X).
length_count: along-track pings (length resolution, Y). Defaults to beam_count.
Each scene has exactly beam_count × length_count sonar returns (first-hit casting).
"""
count = int(count)
if count < 1:
@@ -688,6 +622,7 @@ def generate_dataset(
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")
@@ -723,11 +658,19 @@ def generate_dataset(
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=object_scale,
object_scale=scene_scale,
beam_count=beam_count,
length_count=length_count,
)
@@ -743,6 +686,7 @@ def generate_dataset(
"hasObject": scene["hasObject"],
"pointCount": scene["pointCount"],
"objectPointCount": scene["objectPointCount"],
"objectScale": scene_scale,
"files": paths,
}
written.append(entry)
@@ -777,6 +721,7 @@ def generate_dataset(
"lengthCount": length_count,
"objectName": object_name,
"objectScale": object_scale,
"objectScaleIsMax": object_scale_is_max,
"objectVertexCount": len(template),
"classLabels": {"0": "background", "1": "object"},
"stats": stats,