mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
111 lines
3.0 KiB
Python
111 lines
3.0 KiB
Python
# /// script
|
|
# requires-python = ">=3.13"
|
|
# dependencies = [
|
|
# "marimo>=0.21.1",
|
|
# "pyzmq>=27.1.0",
|
|
# ]
|
|
# ///
|
|
|
|
# pyright: reportMissingImports=false, reportMissingParameterType=false, reportUnknownParameterType=false, reportUnknownVariableType=false, reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnusedParameter=false, reportUnusedCallResult=false, reportUnusedExpression=false
|
|
|
|
import marimo
|
|
|
|
__generated_with = "0.22.4"
|
|
app = marimo.App(width="medium", app_title="Test Scene Image Collector")
|
|
|
|
|
|
@app.cell
|
|
def _():
|
|
from pathlib import Path
|
|
|
|
import marimo as mo
|
|
|
|
return Path, mo
|
|
|
|
|
|
@app.cell
|
|
def _(Path, mo):
|
|
SCENES = {
|
|
"skokloster-castle": "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb",
|
|
"apartment_1": "data/scene_datasets/habitat-test-scenes/apartment_1.glb",
|
|
"van-gogh-room": "data/scene_datasets/habitat-test-scenes/van-gogh-room.glb",
|
|
}
|
|
IMAGE_SIZE = 1024
|
|
VIEWS_PER_POINT = 12
|
|
POINTS_PER_SCENE = 5
|
|
SEED = 42
|
|
OUTPUT_DIR = Path("outputs/test_image")
|
|
|
|
mo.md(
|
|
f"## Configuration\n- Scenes: {len(SCENES)}\n- Points/scene: {POINTS_PER_SCENE}\n- Views/point: {VIEWS_PER_POINT}\n- Resolution: {IMAGE_SIZE}x{IMAGE_SIZE}\n- Total images: {len(SCENES) * POINTS_PER_SCENE * VIEWS_PER_POINT}"
|
|
)
|
|
return (
|
|
IMAGE_SIZE,
|
|
OUTPUT_DIR,
|
|
POINTS_PER_SCENE,
|
|
SCENES,
|
|
SEED,
|
|
VIEWS_PER_POINT,
|
|
)
|
|
|
|
|
|
@app.cell
|
|
def _(IMAGE_SIZE, OUTPUT_DIR, POINTS_PER_SCENE, SCENES, SEED, VIEWS_PER_POINT):
|
|
from simulator import collect_scene_images
|
|
|
|
scene_stats = {}
|
|
total_images = 0
|
|
|
|
for scene_name, scene_path in SCENES.items():
|
|
print(f"Collecting {scene_name}...")
|
|
collected = collect_scene_images(
|
|
scene_name=scene_name,
|
|
scene_path=scene_path,
|
|
output_dir=OUTPUT_DIR,
|
|
image_size=IMAGE_SIZE,
|
|
views_per_point=VIEWS_PER_POINT,
|
|
points_per_scene=POINTS_PER_SCENE,
|
|
seed=SEED,
|
|
)
|
|
scene_stats[scene_name] = collected
|
|
total_images += collected
|
|
print(f"Collected {collected} images for {scene_name}")
|
|
return scene_stats, total_images
|
|
|
|
|
|
@app.cell
|
|
def _(OUTPUT_DIR, mo, scene_stats, total_images):
|
|
mo.md(f"""
|
|
## Summary\n- Output directory: `{OUTPUT_DIR}`\n- Total images: {total_images}\n- Per-scene: {scene_stats}
|
|
""")
|
|
return
|
|
|
|
|
|
@app.cell
|
|
def _(OUTPUT_DIR, SCENES, mo):
|
|
from PIL import Image
|
|
|
|
preview_grid = []
|
|
for _scene_name in SCENES:
|
|
sample_path = OUTPUT_DIR / _scene_name / "000" / "view_000.png"
|
|
if not sample_path.exists():
|
|
continue
|
|
preview_grid.append(
|
|
mo.vstack(
|
|
[mo.md(f"**{_scene_name}**"), mo.image(Image.open(sample_path))],
|
|
align="center",
|
|
)
|
|
)
|
|
|
|
preview = (
|
|
mo.vstack(preview_grid, align="center")
|
|
if preview_grid
|
|
else mo.md("No preview images yet.")
|
|
)
|
|
preview
|
|
return
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run()
|