Сохраняемся

This commit is contained in:
2026-07-24 09:45:32 +03:00
parent 6cfc16e53c
commit 0614ea384a
14 changed files with 1592 additions and 142 deletions
+350 -23
View File
@@ -20,6 +20,7 @@ from scene_generator import (
apply_transform,
export_npy_float64,
export_obj,
generate_pipe,
parse_obj_labeled_points,
parse_obj_points,
)
@@ -37,6 +38,24 @@ TOTAL_FULL_SCENES = sum(n for _, n in AREA_LAYOUT) # 500
VISIBILITY_TIERS = ("nearly_hidden", "partial", "visible")
# Built-in object models for the dataset generator (no upload required).
OBJECT_PRESETS: dict[str, dict[str, Any]] = {
"pipe": {
"id": "pipe",
"label": "Трубопровод",
"filename": "pipeline.obj",
# Length is scene-sized at placement time (border→border); params kept for catalog.
"params": {
"length": 2.6,
"radius": 0.22,
"axis": "y",
"count": 4000,
"noise": 0.005,
"seed": 1,
},
},
}
# ---------------------------------------------------------------------------
# Area naming
@@ -62,9 +81,37 @@ def scene_index_to_area_name(index: int) -> tuple[int, int, str]:
# ---------------------------------------------------------------------------
# Target object from user .obj
# Target object from user .obj or built-in preset
# ---------------------------------------------------------------------------
def list_object_presets() -> list[dict[str, Any]]:
"""Public preset catalog for the dataset UI."""
return [
{
"id": meta["id"],
"label": meta["label"],
"filename": meta["filename"],
}
for meta in OBJECT_PRESETS.values()
]
def load_preset_object_points(preset_id: str) -> tuple[list[list[float]], str]:
"""Build normalized object points for a built-in preset. Returns (points, filename)."""
key = str(preset_id or "").strip().lower()
meta = OBJECT_PRESETS.get(key)
if meta is None:
known = ", ".join(sorted(OBJECT_PRESETS)) or "(none)"
raise ValueError(f"Unknown object preset '{preset_id}'. Known: {known}")
if key == "pipe":
points = generate_pipe(dict(meta["params"]))
else:
raise ValueError(f"Preset '{preset_id}' has no generator.")
if len(points) < 3:
raise ValueError(f"Preset '{preset_id}' produced too few points.")
return normalize_object_points(points), str(meta["filename"])
def normalize_object_points(points: list[list[float]]) -> list[list[float]]:
xs = [p[0] for p in points]
ys = [p[1] for p in points]
@@ -286,6 +333,169 @@ def _place_object_in_scene(
return world, info
def _terrain_unevenness(meta: dict[str, Any]) -> float:
"""Summarize seafloor relief strength as a ~0…1 scalar for silt scaling."""
amp = float(meta.get("amplitude", 0.0))
hills = meta.get("hills") or []
valleys = meta.get("valleys") or []
bumps = meta.get("bumps") or []
hill_peak = max((float(h[2]) for h in hills), default=0.0)
valley_peak = max((float(v[2]) for v in valleys), default=0.0)
bump_peak = max((float(b[2]) for b in bumps), default=0.0)
score = (
amp / 0.35
+ 0.5 * (hill_peak / 0.7)
+ 0.3 * (valley_peak / 0.55)
+ 0.25 * (bump_peak / 0.35)
+ 0.04 * len(hills)
+ 0.02 * len(bumps)
)
return max(0.0, min(1.0, score / 2.2))
def _place_pipeline_in_scene(
rng: random.Random,
meta: dict[str, Any],
visibility: str,
object_scale: float = 1.0,
) -> tuple[list[list[float]], dict[str, Any]]:
"""Lay a straight pipeline from scan border to border; may be silted over.
Length spans the full survey rectangle along X or Y (small yaw allowed).
For ``visible`` there is no silt. For partial / nearly_hidden, randomly
one silt mound or several (count from visibility and seafloor unevenness).
"""
size_x = float(meta["sizeX"])
size_y = float(meta["sizeY"])
min_span = min(size_x, size_y)
base_scale = max(0.01, float(object_scale))
uneven = _terrain_unevenness(meta)
# Diameter scales with scene and UI objectScale; length always edge-to-edge.
radius = max(0.05, min_span * 0.028 * base_scale)
radius = min(radius, min_span * 0.12)
along_x = rng.random() < 0.5
yaw = rng.uniform(-math.radians(10), math.radians(10))
cos_y = math.cos(yaw)
sin_y = math.sin(yaw)
if along_x:
length = size_x / max(abs(cos_y), 0.88) * 1.04
cx = 0.0
cy = rng.uniform(-size_y * 0.12, size_y * 0.12)
ux, uy = cos_y, sin_y
nx, ny = -sin_y, cos_y
else:
length = size_y / max(abs(cos_y), 0.88) * 1.04
cx = rng.uniform(-size_x * 0.12, size_x * 0.12)
cy = 0.0
ux, uy = -sin_y, cos_y
nx, ny = -cos_y, -sin_y
def centerline_xy(t_norm: float) -> tuple[float, float]:
t = t_norm * length
return cx + ux * t, cy + uy * t
# Ось — прямая в 3D: одна высота для всего цилиндра (не повторяет рельеф).
n_along = max(96, int(max(int(meta["beamCount"]), int(meta["lengthCount"])) * 3))
n_circ = 28
floor_samples = [
_height_at(*centerline_xy(ia / max(n_along - 1, 1) - 0.5), meta)
for ia in range(n_along)
]
# Положить трубу примерно на средний уровень дна, не изгибая ось.
axis_z = (sum(floor_samples) / max(len(floor_samples), 1)) + radius
world: list[list[float]] = []
for ia in range(n_along):
t_norm = ia / max(n_along - 1, 1) - 0.5 # [-0.5, 0.5]
px, py = centerline_xy(t_norm)
for ic in range(n_circ):
ang = (2.0 * math.pi * ic) / n_circ
rr = math.cos(ang) * radius
rz = math.sin(ang) * radius
world.append([px + nx * rr, py + ny * rr, axis_z + rz])
# Ил: только partial / nearly_hidden — случайно один холм или несколько.
silt_bumps: list[tuple[float, float, float, float]] = []
silt_mode = "none"
if visibility in ("nearly_hidden", "partial"):
silt_mode = "several" if rng.random() < 0.5 else "one"
if silt_mode == "one":
n_mounds = 1
elif visibility == "nearly_hidden":
n_mounds = rng.randint(2, 5) + (1 if uneven > 0.55 else 0)
else:
n_mounds = rng.randint(2, 4) + (1 if uneven > 0.65 else 0)
if visibility == "nearly_hidden":
span = 0.22 + 0.28 * uneven
cover = 0.85 + 0.55 * uneven
else:
span = 0.14 + 0.22 * uneven
cover = 0.4 + 0.45 * uneven
# Один холм — шире/выше; несколько — компактнее и разнесены.
for i in range(n_mounds):
if n_mounds == 1:
t_norm = rng.uniform(-span * 0.55, span * 0.55)
else:
t_norm = -span + (2.0 * span) * (i + 0.5) / n_mounds
t_norm += rng.uniform(-span * 0.12, span * 0.12)
t_norm = max(-span, min(span, t_norm))
lateral = rng.uniform(-radius * (1.0 + uneven), radius * (1.0 + uneven))
px, py = centerline_xy(t_norm)
sx = px + nx * lateral
sy = py + ny * lateral
local_floor = _height_at(sx, sy, meta)
neighbor = _height_at(sx + radius * 2, sy + radius * 2, meta)
local_relief = abs(local_floor - neighbor) / max(radius, 1e-3)
local_factor = 0.75 + min(0.55, 0.35 * local_relief + 0.25 * uneven)
size_boost = 1.35 if n_mounds == 1 else 1.0
height = radius * cover * local_factor * size_boost * rng.uniform(0.85, 1.2)
mound_r = max(
radius * rng.uniform(2.8, 4.8) * (1.35 if n_mounds == 1 else 1.0),
min_span * rng.uniform(0.05, 0.1) * (0.85 + 0.3 * uneven),
)
silt_bumps.append((sx, sy, height, mound_r))
else:
cover = 0.0
if silt_bumps:
meta.setdefault("bumps", []).extend(silt_bumps)
info = {
"visibility": visibility,
"kind": "pipe",
"transform": {
"x": cx,
"y": cy,
"z": 0.0,
"rx": 0.0,
"ry": 0.0,
"rz": yaw if along_x else (yaw + 0.5 * math.pi),
},
"scale": base_scale,
"objectScale": base_scale,
"radius": radius,
"length": length,
"axis": "x" if along_x else "y",
"hasBend": False,
"bendAmp": 0.0,
"axisZ": round(axis_z, 4),
"burial": 0.0,
"terrainUnevenness": round(uneven, 4),
"siltMode": silt_mode,
"siltCover": round(cover, 4),
"siltBumpCount": len(silt_bumps),
"vertexCount": len(world),
"classLabel": "object",
"classId": 1,
}
return world, info
def _rasterize_object_hits(
world_pts: list[list[float]],
meta: dict[str, Any],
@@ -368,18 +578,72 @@ def _cast_echosounder_returns(
# Balance plan + single scene
# ---------------------------------------------------------------------------
def plan_scene_labels(count: int, seed: int) -> list[str]:
def _allocate_by_weights(total: int, weights: list[float]) -> list[int]:
"""Split ``total`` into integer buckets proportional to non-negative weights."""
n = len(weights)
if total <= 0 or n == 0:
return [0] * n
w = [max(0.0, float(x)) for x in weights]
s = sum(w)
if s <= 0:
out = [0] * n
out[0] = total
return out
raw = [total * wi / s for wi in w]
floors = [int(math.floor(r)) for r in raw]
rem = total - sum(floors)
order = sorted(range(n), key=lambda i: (raw[i] - floors[i], -i), reverse=True)
for i in range(rem):
floors[order[i % n]] += 1
return floors
def _clamp_pct(value: float, *, name: str) -> float:
v = float(value)
if not math.isfinite(v):
raise ValueError(f"{name} must be a finite number")
if v < 0 or v > 100:
raise ValueError(f"{name} must be in [0, 100], got {v}")
return v
def plan_scene_labels(
count: int,
seed: int,
*,
absent_pct: float = 30.0,
nearly_hidden_pct: float = 20.0,
partial_pct: float = 40.0,
visible_pct: float = 40.0,
) -> list[str]:
"""Return visibility label per scene: absent | nearly_hidden | partial | visible.
~50% absent; among present scenes, roughly equal nearly_hidden/partial/visible.
Defaults: 30% absent / 70% with object. Of scenes with an object:
20% nearly_hidden, 40% partial, 40% visible.
Visibility percents are relative to with-object scenes (normalized if needed).
"""
count = max(0, int(count))
absent_pct = _clamp_pct(absent_pct, name="absent_pct")
nearly_hidden_pct = _clamp_pct(nearly_hidden_pct, name="nearly_hidden_pct")
partial_pct = _clamp_pct(partial_pct, name="partial_pct")
visible_pct = _clamp_pct(visible_pct, name="visible_pct")
n_without, n_with = _allocate_by_weights(count, [absent_pct, 100.0 - absent_pct])
tier_weights = [nearly_hidden_pct, partial_pct, visible_pct]
if n_with > 0 and sum(tier_weights) <= 0:
raise ValueError(
"Visibility percents among with-object scenes must sum to > 0 "
"when there are scenes with an object."
)
n_nearly, n_partial, n_visible = _allocate_by_weights(n_with, tier_weights)
labels: list[str] = (
["absent"] * n_without
+ ["nearly_hidden"] * n_nearly
+ ["partial"] * n_partial
+ ["visible"] * n_visible
)
rng = random.Random(int(seed) ^ 0xA5A5_5A5A)
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
@@ -390,6 +654,7 @@ def generate_sonar_scene(
visibility: str = "absent",
object_points: list[list[float]] | None = None,
object_scale: float = 1.0,
object_kind: str | None = None,
beam_count: int = 45,
length_count: int | None = None,
) -> dict[str, Any]:
@@ -397,11 +662,13 @@ def generate_sonar_scene(
visibility in absent|nearly_hidden|partial|visible.
Scene size is always beam_count × length_count returns.
object_kind ``pipe`` places a border-to-border pipeline with optional silt.
"""
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:
kind = (object_kind or "").strip().lower() or None
if visibility != "absent" and kind != "pipe" and not object_points:
raise ValueError("object_points required when visibility is not absent.")
meta = _build_seafloor_meta(rng, beam_count=beam_count, length_count=length_count)
@@ -410,13 +677,21 @@ def generate_sonar_scene(
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,
)
if kind == "pipe":
world, object_info = _place_pipeline_in_scene(
rng,
meta,
visibility,
object_scale=object_scale,
)
else:
world, object_info = _place_object_in_scene(
rng,
meta,
visibility,
object_points,
object_scale=object_scale,
)
object_hits = _rasterize_object_hits(world, meta)
object_info["hitCellCount"] = len(object_hits)
object_info["keptCount"] = len(object_hits)
@@ -533,6 +808,10 @@ def build_run_manifest(
object_vertex_count: int,
stats: dict[str, Any],
written: list[dict[str, Any]],
absent_pct: float = 30.0,
nearly_hidden_pct: float = 20.0,
partial_pct: float = 40.0,
visible_pct: float = 40.0,
) -> dict[str, Any]:
model_filename = _normalize_model_filename(object_name)
return {
@@ -551,6 +830,10 @@ def build_run_manifest(
"beamCount": int(beam_count),
"lengthCount": int(length_count),
"objectName": model_filename,
"absentPct": float(absent_pct),
"nearlyHiddenPct": float(nearly_hidden_pct),
"partialPct": float(partial_pct),
"visiblePct": float(visible_pct),
},
"objectVertexCount": int(object_vertex_count),
"stats": stats,
@@ -850,12 +1133,17 @@ def iter_generate_dataset(
count: int = 5,
seed: int = 42,
output_dir: str | Path = "sonar_dataset",
object_points: list[list[float]],
object_points: list[list[float]] | None = None,
object_name: str | None = None,
object_kind: str | None = None,
object_scale: float = 1.0,
object_scale_is_max: bool = False,
beam_count: int = 45,
length_count: int | None = None,
absent_pct: float = 30.0,
nearly_hidden_pct: float = 20.0,
partial_pct: float = 40.0,
visible_pct: float = 40.0,
):
"""Yield NDJSON-friendly progress events, then a final ``done`` payload.
@@ -869,8 +1157,13 @@ def iter_generate_dataset(
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.")
kind = (object_kind or "").strip().lower() or None
if kind == "pipe":
template: list[list[float]] = []
else:
if not object_points or len(object_points) < 3:
raise ValueError("A valid .obj model with at least 3 vertices is required.")
template = normalize_object_points(object_points)
object_scale = float(object_scale)
if object_scale <= 0:
raise ValueError("object_scale must be > 0")
@@ -890,13 +1183,24 @@ def iter_generate_dataset(
if length_count > 1024:
raise ValueError("length_count (Длина) must be <= 1024")
absent_pct = _clamp_pct(absent_pct, name="absent_pct")
nearly_hidden_pct = _clamp_pct(nearly_hidden_pct, name="nearly_hidden_pct")
partial_pct = _clamp_pct(partial_pct, name="partial_pct")
visible_pct = _clamp_pct(visible_pct, name="visible_pct")
object_name = _normalize_model_filename(object_name)
base = resolve_output_dir(output_dir)
run_dir = make_generation_run_dir(base, object_name)
template = normalize_object_points(object_points)
labels = plan_scene_labels(count, seed)
labels = plan_scene_labels(
count,
seed,
absent_pct=absent_pct,
nearly_hidden_pct=nearly_hidden_pct,
partial_pct=partial_pct,
visible_pct=visible_pct,
)
written: list[dict[str, Any]] = []
stats = {
"total": count,
@@ -911,6 +1215,7 @@ def iter_generate_dataset(
preview_points: list[list[float]] | None = None
preview_stem: str | None = None
preview_has_object = False
object_vertex_count = len(template)
yield {
"type": "start",
@@ -936,9 +1241,12 @@ def iter_generate_dataset(
visibility=visibility,
object_points=template,
object_scale=scene_scale,
object_kind=kind,
beam_count=beam_count,
length_count=length_count,
)
if scene.get("object") and scene["object"].get("vertexCount"):
object_vertex_count = int(scene["object"]["vertexCount"])
area, scene_no, stem = scene_index_to_area_name(i)
paths = write_scene_files(scene, run_dir, stem)
@@ -995,9 +1303,14 @@ def iter_generate_dataset(
"beamCount": beam_count,
"lengthCount": length_count,
"objectName": object_name,
"objectKind": kind,
"objectScale": object_scale,
"objectScaleIsMax": object_scale_is_max,
"objectVertexCount": len(template),
"objectVertexCount": object_vertex_count,
"absentPct": absent_pct,
"nearlyHiddenPct": nearly_hidden_pct,
"partialPct": partial_pct,
"visiblePct": visible_pct,
"classLabels": {"0": "background", "1": "object"},
"stats": stats,
"written": written,
@@ -1017,6 +1330,10 @@ def iter_generate_dataset(
object_vertex_count=len(template),
stats=stats,
written=written,
absent_pct=absent_pct,
nearly_hidden_pct=nearly_hidden_pct,
partial_pct=partial_pct,
visible_pct=visible_pct,
)
write_run_manifest(run_dir, manifest)
result["settingsPath"] = str(run_dir / DATASET_RUN_FILENAME)
@@ -1028,12 +1345,17 @@ def generate_dataset(
count: int = 5,
seed: int = 42,
output_dir: str | Path = "sonar_dataset",
object_points: list[list[float]],
object_points: list[list[float]] | None = None,
object_name: str | None = None,
object_kind: str | None = None,
object_scale: float = 1.0,
object_scale_is_max: bool = False,
beam_count: int = 45,
length_count: int | None = None,
absent_pct: float = 30.0,
nearly_hidden_pct: float = 20.0,
partial_pct: float = 40.0,
visible_pct: float = 40.0,
) -> dict[str, Any]:
"""Generate `count` unique scenes into a new timestamped run folder under output_dir."""
result: dict[str, Any] | None = None
@@ -1043,10 +1365,15 @@ def generate_dataset(
output_dir=output_dir,
object_points=object_points,
object_name=object_name,
object_kind=object_kind,
object_scale=object_scale,
object_scale_is_max=object_scale_is_max,
beam_count=beam_count,
length_count=length_count,
absent_pct=absent_pct,
nearly_hidden_pct=nearly_hidden_pct,
partial_pct=partial_pct,
visible_pct=visible_pct,
):
if event.get("type") == "done":
result = event["result"]