Сохранять каждый запуск датасета в отдельную папку и показывать прогресс генерации.

Файлы пишутся в sonar_dataset/дата-время-модель без удаления прошлых запусков; кнопка «Генерация…» заполняется по мере записи сцен через NDJSON-стрим.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
2026-07-17 15:30:07 +03:00
co-authored by Cursor
parent 643d834ba7
commit 8b248e1d54
5 changed files with 258 additions and 42 deletions
+92 -14
View File
@@ -474,6 +474,35 @@ def resolve_output_dir(output_dir: str | Path = "sonar_dataset") -> Path:
return out
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-<object_stem>.
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 _downsample_points(points: list[list[float]], max_points: int) -> list[list[float]]:
max_points = max(100, int(max_points))
if len(points) <= max_points:
@@ -588,7 +617,7 @@ def write_scene_files(
return {"npy": str(npy_path), "obj": str(obj_path), "stem": stem}
def generate_dataset(
def iter_generate_dataset(
*,
count: int = 5,
seed: int = 42,
@@ -599,16 +628,13 @@ def generate_dataset(
object_scale_is_max: bool = False,
beam_count: int = 45,
length_count: int | None = None,
) -> dict[str, Any]:
"""Generate `count` unique scenes into output_dir with Area_X naming.
):
"""Yield NDJSON-friendly progress events, then a final ``done`` payload.
object_points: normalized template vertices from user .obj (class 1 = object).
object_scale: relative size multiplier vs unit-normalized mesh (1.0 = default).
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).
Events:
{"type":"start","total":N,"outputDir":"...","runName":"..."}
{"type":"progress","current":k,"total":N,"entry":{...}}
{"type":"done","result":{...}}
"""
count = int(count)
if count < 1:
@@ -636,7 +662,8 @@ def generate_dataset(
if length_count > 1024:
raise ValueError("length_count (Длина) must be <= 1024")
out = resolve_output_dir(output_dir)
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)
@@ -655,6 +682,14 @@ def generate_dataset(
preview_stem: str | None = None
preview_has_object = False
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
@@ -675,7 +710,7 @@ def generate_dataset(
length_count=length_count,
)
area, scene_no, stem = scene_index_to_area_name(i)
paths = write_scene_files(scene, out, stem)
paths = write_scene_files(scene, run_dir, stem)
entry = {
"index": i,
@@ -702,6 +737,14 @@ def generate_dataset(
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)
@@ -713,8 +756,10 @@ def generate_dataset(
"classLabels": {"0": "background", "1": "object"},
}
return {
"outputDir": str(out),
result = {
"outputDir": str(run_dir),
"baseDir": str(base),
"runName": run_dir.name,
"count": count,
"seed": int(seed),
"beamCount": beam_count,
@@ -728,3 +773,36 @@ def generate_dataset(
"written": written,
"preview": preview,
}
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]],
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]:
"""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_scale=object_scale,
object_scale_is_max=object_scale_is_max,
beam_count=beam_count,
length_count=length_count,
):
if event.get("type") == "done":
result = event["result"]
if result is None:
raise RuntimeError("Dataset generation produced no result.")
return result
+27 -16
View File
@@ -11,14 +11,14 @@ from typing import Any
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, Response
from fastapi.responses import FileResponse, Response, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from builtin_presets import BUILTIN_PRESETS, get_builtin_preset
from demo_generator import DEMO_SURFACE_TYPES, demo_payload
from pipeline_insights import compute_insights
from dataset_generator import generate_dataset, load_object_points_from_obj_text, load_scene_preview
from dataset_generator import iter_generate_dataset, load_object_points_from_obj_text, load_scene_preview
from scene_generator import (
catalog_payload as generator_catalog_payload,
export_npy_float64,
@@ -429,7 +429,7 @@ async def dataset_generate(
beamCount: int = Form(45),
lengthCount: int | None = Form(None),
model: UploadFile = File(...),
) -> dict[str, Any]:
) -> StreamingResponse:
filename = (model.filename or "").strip()
if not filename.lower().endswith(".obj"):
raise HTTPException(status_code=400, detail="Upload a .obj 3D model file.")
@@ -437,21 +437,32 @@ async def dataset_generate(
raw = await model.read()
text = raw.decode("utf-8", errors="ignore")
object_points = load_object_points_from_obj_text(text)
return generate_dataset(
count=count,
seed=seed,
output_dir=outputDir or "sonar_dataset",
object_points=object_points,
object_name=filename,
object_scale=objectScale,
object_scale_is_max=objectScaleIsMax,
beam_count=beamCount,
length_count=lengthCount,
)
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
except OSError as exc:
raise HTTPException(status_code=500, detail=f"Failed to write dataset: {exc}") from exc
def event_stream():
try:
for event in iter_generate_dataset(
count=count,
seed=seed,
output_dir=outputDir or "sonar_dataset",
object_points=object_points,
object_name=filename,
object_scale=objectScale,
object_scale_is_max=objectScaleIsMax,
beam_count=beamCount,
length_count=lengthCount,
):
yield json.dumps(event, ensure_ascii=False) + "\n"
except ValueError as exc:
yield json.dumps({"type": "error", "detail": str(exc)}, ensure_ascii=False) + "\n"
except OSError as exc:
yield json.dumps(
{"type": "error", "detail": f"Failed to write dataset: {exc}"},
ensure_ascii=False,
) + "\n"
return StreamingResponse(event_stream(), media_type="application/x-ndjson")
@app.post("/api/dataset/preview")