Spaces:
Sleeping
Sleeping
File size: 3,754 Bytes
9468cdd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """
FastAPI server for the OpenX Rerun Viewer HF Space.
Runs the selection UI and serves .rrd files.
When running locally, also starts a local Rerun web viewer to avoid mixed-content issues.
"""
import json
import os
import subprocess
import sys
import time
from pathlib import Path
from fastapi import FastAPI
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
app = FastAPI(title="OpenX Rerun Viewer")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["GET", "HEAD", "OPTIONS"],
allow_headers=["*"],
expose_headers=[
"Content-Length", "Content-Type", "Accept-Ranges",
"Cross-Origin-Resource-Policy",
],
)
DATA_DIR = Path(__file__).parent / "data"
MANIFEST_PATH = DATA_DIR / "manifest.json"
IS_LOCAL = "SPACE_ID" not in os.environ and "HF_SPACE" not in os.environ
def load_manifest():
if MANIFEST_PATH.exists():
return json.loads(MANIFEST_PATH.read_text())
return []
@app.get("/", response_class=HTMLResponse)
async def index():
html = (Path(__file__).parent / "templates" / "index.html").read_text()
# Inject the viewer base URL: use local viewer for dev, app.rerun.io for production
if IS_LOCAL:
viewer_base = "http://localhost:9090"
else:
viewer_base = "https://app.rerun.io/version/0.26.2"
html = html.replace("{{VIEWER_BASE}}", viewer_base)
return HTMLResponse(html)
@app.get("/api/datasets")
async def list_datasets():
manifest = load_manifest()
return [
{
"slug": m["slug"],
"robot_type": m["robot_type"],
"fps": m.get("fps", 0),
"camera_keys": m.get("camera_keys", []),
"camera_labels": m.get("camera_labels", []),
"num_episodes_total": m.get("num_episodes_total", 0),
"state_dim": m.get("state_dim", 0),
"action_dim": m.get("action_dim", 0),
"episodes_available": [
{"index": ep["index"], "frames": ep.get("frames", "?")}
for ep in m.get("episodes_available", [])
],
}
for m in manifest
]
@app.get("/api/datasets/{slug}/episodes")
async def list_episodes(slug: str):
manifest = load_manifest()
for m in manifest:
if m["slug"] == slug:
return m["episodes_available"]
return []
@app.api_route("/data/{slug}/episode_{index:str}.rrd", methods=["GET", "HEAD"])
async def serve_rrd(slug: str, index: str):
filepath = DATA_DIR / slug / f"episode_{index}.rrd"
if not filepath.exists():
return HTMLResponse(status_code=404, content="Episode not found")
return FileResponse(
filepath,
media_type="application/octet-stream",
headers={
"Accept-Ranges": "bytes",
"Access-Control-Allow-Origin": "*",
"Cross-Origin-Resource-Policy": "cross-origin",
"Cache-Control": "public, max-age=3600",
},
)
static_dir = Path(__file__).parent / "static"
if static_dir.exists():
app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")
def start_local_viewer():
"""Start Rerun web viewer on port 9090 for local development."""
import rerun as rr
rr.serve_web_viewer(web_port=9090, open_browser=False)
# serve_web_viewer returns immediately after starting the server in a thread
if __name__ == "__main__":
import uvicorn
port = int(os.environ.get("PORT", 8000))
if IS_LOCAL:
print("Local mode: starting Rerun web viewer on http://localhost:9090")
start_local_viewer()
time.sleep(0.5)
uvicorn.run(app, host="0.0.0.0", port=port)
|