Сохранять каждый запуск датасета в отдельную папку и показывать прогресс генерации.
Файлы пишутся в sonar_dataset/дата-время-модель без удаления прошлых запусков; кнопка «Генерация…» заполняется по мере записи сцен через NDJSON-стрим. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -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
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user