"""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, generate_pipe, 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") # 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 # --------------------------------------------------------------------------- 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 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] 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, mle_mode: bool = False, auv_depth: float | None = None, swath_angle_deg: float = 90.0, relief_scale_pct: float = 100.0, ) -> dict[str, Any]: """Build continuous seafloor heightfield parameters (no point cloud yet). ``relief_scale_pct`` scales random unevenness relative to scene size (100 = baseline, clearly visible relief; higher = larger features). """ beams = max(1, int(beam_count)) length_points = beams if length_count is None else max(1, int(length_count)) swath_deg = max(5.0, min(170.0, float(swath_angle_deg))) relief_pct = max(5.0, min(500.0, float(relief_scale_pct))) if mle_mode and auv_depth is not None: # Swath footprint from altitude above flat datum (like МЛЭ). depth = max(0.3, float(auv_depth)) half_w = depth * math.tan(math.radians(swath_deg) * 0.5) size_x = max(4.0, 2.0 * half_w * 1.08) dx = size_x / max(beams - 1, 1) size_y = max(4.0, dx * max(length_points - 1, 1) * rng.uniform(0.95, 1.12)) else: size_x = rng.uniform(8.0, 16.0) size_y = rng.uniform(8.0, 16.0) base_z = rng.uniform(-1.2, -0.2) # Feature size tied to scene so relief stays visible in the point cloud. # At 100%: characteristic scale ≈ 28% of the shorter scene side. scene_ref = max(4.0, min(size_x, size_y)) feature_scale = scene_ref * (relief_pct / 100.0) * 0.28 feature_scale = max(0.35, feature_scale) max_rad = scene_ref * 0.42 # Background undulation — several waves across the patch. amplitude = feature_scale * rng.uniform(0.12, 0.28) frequency = (2.0 * math.pi / scene_ref) * rng.uniform(1.8, 3.6) n_hills = rng.randint(2, 5) hills = [] for _ in range(n_hills): hills.append( ( rng.uniform(-size_x * 0.38, size_x * 0.38), rng.uniform(-size_y * 0.38, size_y * 0.38), feature_scale * rng.uniform(0.45, 1.05), min(max_rad, feature_scale * rng.uniform(0.75, 1.55)), ) ) n_valleys = rng.randint(1, 4) valleys = [] for _ in range(n_valleys): valleys.append( ( rng.uniform(-size_x * 0.38, size_x * 0.38), rng.uniform(-size_y * 0.38, size_y * 0.38), feature_scale * rng.uniform(0.3, 0.75), min(max_rad, feature_scale * rng.uniform(0.65, 1.4)), ) ) # More medium bumps at lower scale; fewer but larger when scale is high. density = max(0.45, min(1.8, 120.0 / max(relief_pct, 20.0))) n_bumps = int(round(rng.uniform(10, 20) * density)) n_bumps = max(8, min(40, n_bumps)) bumps = [] for _ in range(n_bumps): bumps.append( ( rng.uniform(-size_x * 0.46, size_x * 0.46), rng.uniform(-size_y * 0.46, size_y * 0.46), feature_scale * rng.uniform(0.12, 0.45), min(max_rad * 0.7, feature_scale * rng.uniform(0.25, 0.85)), ) ) 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), "mleMode": bool(mle_mode), "auvDepth": float(auv_depth) if auv_depth is not None else None, "swathAngleDeg": swath_deg, "reliefScalePct": relief_pct, "featureScale": feature_scale, } 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 _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], ) -> 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 def _build_object_height_field( world_pts: list[list[float]], meta: dict[str, Any], ) -> dict[str, Any] | None: """Dense XY max-Z map of object surface for angled multibeam probes.""" if not world_pts: return None beams = int(meta["beamCount"]) length_points = int(meta["lengthCount"]) nx = max(32, min(256, beams * 3)) ny = max(32, min(256, length_points * 3)) size_x = float(meta["sizeX"]) size_y = float(meta["sizeY"]) half_x = size_x * 0.5 half_y = size_y * 0.5 cell_w = size_x / max(nx - 1, 1) cell_h = size_y / max(ny - 1, 1) splat_rx = cell_w * 1.1 splat_ry = cell_h * 1.1 field: list[list[float | None]] = [[None for _ in range(nx)] for _ in range(ny)] for px, py, pz in world_pts: fx = 0.0 if nx == 1 else (px + half_x) / size_x * (nx - 1) fy = 0.0 if ny == 1 else (py + half_y) / size_y * (ny - 1) 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 >= nx or yi >= ny: continue cx = -half_x if nx == 1 else (-half_x + size_x * xi / (nx - 1)) cy = -half_y if ny == 1 else (-half_y + size_y * yi / (ny - 1)) if abs(px - cx) > splat_rx or abs(py - cy) > splat_ry: continue floor_z = _height_at(cx, cy, meta) if pz < floor_z - 0.01: continue prev = field[yi][xi] if prev is None or pz > prev: field[yi][xi] = pz return { "nx": nx, "ny": ny, "sizeX": size_x, "sizeY": size_y, "field": field, } def _sample_object_height_field( field_info: dict[str, Any] | None, x: float, y: float, ) -> float | None: if not field_info: return None nx = int(field_info["nx"]) ny = int(field_info["ny"]) size_x = float(field_info["sizeX"]) size_y = float(field_info["sizeY"]) half_x = size_x * 0.5 half_y = size_y * 0.5 u = (x + half_x) / max(size_x, 1e-6) v = (y + half_y) / max(size_y, 1e-6) if u < -0.02 or u > 1.02 or v < -0.02 or v > 1.02: return None u = min(1.0, max(0.0, u)) v = min(1.0, max(0.0, v)) fx = u * (nx - 1) fy = v * (ny - 1) i0 = int(math.floor(fx)) j0 = int(math.floor(fy)) i1 = min(nx - 1, i0 + 1) j1 = min(ny - 1, j0 + 1) vals = [ field_info["field"][j0][i0], field_info["field"][j0][i1], field_info["field"][j1][i0], field_info["field"][j1][i1], ] present = [z for z in vals if z is not None] if not present: return None return max(present) def _cast_mle_fan_returns( rng: random.Random, meta: dict[str, Any], object_field: dict[str, Any] | None, *, auv_depth: float, swath_angle_deg: float = 90.0, ) -> tuple[list[list[float]], int, int]: """Multibeam fan: beamCount rays × lengthCount pings from constant altitude. AUV height is measured from flat seafloor datum (baseZ), not terrain-following. """ beams = int(meta["beamCount"]) length_points = int(meta["lengthCount"]) noise = float(meta.get("noise", 0.01)) base_z = float(meta["baseZ"]) depth = max(0.3, float(auv_depth)) origin_z = base_z + depth swath = math.radians(max(5.0, min(170.0, float(swath_angle_deg)))) size_y = float(meta["sizeY"]) half_y = size_y * 0.5 # Reach outer beams with margin past the swath edge. max_range = max(depth / max(math.cos(swath * 0.5), 0.12) * 1.35, depth * 2.0, 4.0) step = max(0.08, min(0.45, max_range / 220.0)) max_steps = int(math.ceil(max_range / step)) + 2 rows: list[list[float]] = [] object_count = 0 background_count = 0 for yi in range(length_points): y = -half_y if length_points == 1 else (-half_y + size_y * yi / (length_points - 1)) ox, oy, oz = 0.0, y, origin_z for xi in range(beams): t = 0.5 if beams == 1 else xi / (beams - 1) angle = -swath * 0.5 + swath * t dx = math.sin(angle) dy = 0.0 dz = -math.cos(angle) hit_x, hit_y, hit_z = ox, oy, base_z is_object = False px = ox + dx * step * 0.5 py = oy + dy * step * 0.5 pz = oz + dz * step * 0.5 for _ in range(max_steps): dist = math.hypot(px - ox, py - oy, pz - oz) if dist > max_range: break floor_z = _height_at(px, py, meta) obj_z = _sample_object_height_field(object_field, px, py) if obj_z is not None and pz <= obj_z and obj_z >= floor_z - 0.01: hit_x, hit_y, hit_z = px, py, obj_z is_object = True break if pz <= floor_z: hit_x, hit_y, hit_z = px, py, floor_z break px += dx * step py += dy * step pz += dz * step else: # No intersection within range — project to flat nadir estimate. hit_x = ox + dx * max_range hit_y = oy + dy * max_range hit_z = _height_at(hit_x, hit_y, meta) z = hit_z + rng.uniform(-noise, noise) if is_object: cls = 1.0 object_count += 1 else: cls = 0.0 background_count += 1 rows.append([float(hit_x), float(hit_y), float(z), 0.0, 0.0, 0.0, cls]) return rows, object_count, background_count # --------------------------------------------------------------------------- # Balance plan + single scene # --------------------------------------------------------------------------- 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. 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) 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, object_kind: str | None = None, beam_count: int = 45, length_count: int | None = None, mle_mode: bool = False, auv_depth: float | None = None, swath_angle_deg: float = 90.0, relief_scale_pct: float = 100.0, ) -> 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. object_kind ``pipe`` places a border-to-border pipeline with optional silt. When ``mle_mode`` is True, returns are cast as a multibeam fan from constant altitude ``auv_depth`` above the flat seafloor datum (baseZ). """ rng = random.Random(int(seed)) if visibility not in ("absent",) + VISIBILITY_TIERS: raise ValueError(f"Unknown visibility: {visibility}") 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.") use_mle = bool(mle_mode) depth = max(0.3, float(auv_depth)) if auv_depth is not None else None if use_mle and depth is None: depth = 2.5 meta = _build_seafloor_meta( rng, beam_count=beam_count, length_count=length_count, mle_mode=use_mle, auv_depth=depth, swath_angle_deg=swath_angle_deg, relief_scale_pct=relief_scale_pct, ) expected = int(meta["beamCount"]) * int(meta["lengthCount"]) object_info: dict[str, Any] | None = None object_hits: dict[tuple[int, int], float] | None = None object_field: dict[str, Any] | None = None world: list[list[float]] | None = None if visibility != "absent": 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, ) if use_mle: object_field = _build_object_height_field(world, meta) hit_cells = 0 if object_field: for row in object_field["field"]: hit_cells += sum(1 for z in row if z is not None) object_info["hitCellCount"] = hit_cells object_info["keptCount"] = hit_cells else: object_hits = _rasterize_object_hits(world, meta) object_info["hitCellCount"] = len(object_hits) object_info["keptCount"] = len(object_hits) object_info["requestedCount"] = expected if use_mle: rows, object_point_count, background_point_count = _cast_mle_fan_returns( rng, meta, object_field, auv_depth=float(depth), swath_angle_deg=float(meta.get("swathAngleDeg", swath_angle_deg)), ) else: 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, "auvDepth": float(depth) if use_mle else None, "mleMode": use_mle, "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, "mleMode": use_mle, "auvDepth": float(depth) if use_mle else None, "swathAngleDeg": meta.get("swathAngleDeg"), "reliefScalePct": meta.get("reliefScalePct"), "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-. 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]], absent_pct: float = 30.0, nearly_hidden_pct: float = 20.0, partial_pct: float = 40.0, visible_pct: float = 40.0, mle_mode: bool = False, auv_depth_min: float | None = None, auv_depth_max: float | None = None, relief_scale_pct: float = 100.0, ) -> dict[str, Any]: model_filename = _normalize_model_filename(object_name) settings: dict[str, Any] = { "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, "absentPct": float(absent_pct), "nearlyHiddenPct": float(nearly_hidden_pct), "partialPct": float(partial_pct), "visiblePct": float(visible_pct), "mleMode": bool(mle_mode), "reliefScalePct": float(relief_scale_pct), } if mle_mode: settings["auvDepthMin"] = float(auv_depth_min) if auv_depth_min is not None else None settings["auvDepthMax"] = float(auv_depth_max) if auv_depth_max is not None else None 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": settings, "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(" 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]] | 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, mle_mode: bool = False, auv_depth_min: float = 2.0, auv_depth_max: float = 8.0, relief_scale_pct: float = 100.0, ): """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") 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") 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") use_mle = bool(mle_mode) depth_lo = max(0.3, float(auv_depth_min)) depth_hi = max(0.3, float(auv_depth_max)) if depth_hi < depth_lo: depth_lo, depth_hi = depth_hi, depth_lo if use_mle and (not math.isfinite(depth_lo) or not math.isfinite(depth_hi)): raise ValueError("Диапазон высот must be finite numbers") relief_pct = max(5.0, min(500.0, float(relief_scale_pct))) if not math.isfinite(relief_pct): raise ValueError("relief_scale_pct (Размер неровностей) must be finite") 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) 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, "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 object_vertex_count = len(template) 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_depth = scene_rng.uniform(depth_lo, depth_hi) if use_mle else None scene = generate_sonar_scene( seed=scene_seed, visibility=visibility, object_points=template, object_scale=scene_scale, object_kind=kind, beam_count=beam_count, length_count=length_count, mle_mode=use_mle, auv_depth=scene_depth, relief_scale_pct=relief_pct, ) 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) entry = { "index": i, "area": area, "scene": scene_no, "stem": stem, "visibility": visibility, "hasObject": scene["hasObject"], "pointCount": scene["pointCount"], "objectPointCount": scene["objectPointCount"], "objectScale": scene_scale, "auvDepth": scene.get("auvDepth"), "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, "mleMode": use_mle, "auvDepthMin": depth_lo if use_mle else None, "auvDepthMax": depth_hi if use_mle else None, "reliefScalePct": relief_pct, "objectName": object_name, "objectKind": kind, "objectScale": object_scale, "objectScaleIsMax": object_scale_is_max, "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, "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, absent_pct=absent_pct, nearly_hidden_pct=nearly_hidden_pct, partial_pct=partial_pct, visible_pct=visible_pct, mle_mode=use_mle, auv_depth_min=depth_lo if use_mle else None, auv_depth_max=depth_hi if use_mle else None, relief_scale_pct=relief_pct, ) 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]] | 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, mle_mode: bool = False, auv_depth_min: float = 2.0, auv_depth_max: float = 8.0, relief_scale_pct: float = 100.0, ) -> 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_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, mle_mode=mle_mode, auv_depth_min=auv_depth_min, auv_depth_max=auv_depth_max, relief_scale_pct=relief_scale_pct, ): if event.get("type") == "done": result = event["result"] if result is None: raise RuntimeError("Dataset generation produced no result.") return result