asus
5c7c122105
This commit includes major additions and updates to the frontend and backend architectures, introducing new dataset management, live counting features, batch processing, and triage logic. Includes new UI pages, components, and API routes.
86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
"""FastAPI server for the dataset enrichment platform.
|
|
|
|
Run it with: .venv/bin/uvicorn backend.main:app --port 8000
|
|
or: docker compose up
|
|
|
|
Routes live in `backend/api/`, one module per domain; this file only wires them
|
|
together and owns startup.
|
|
"""
|
|
|
|
import os
|
|
import shutil
|
|
from contextlib import asynccontextmanager
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
from backend import archive_index, config, db, jobs
|
|
from backend.api import (batches, counting_bench as counting_bench_routes,
|
|
datasets as dataset_routes, jobs as job_routes,
|
|
live_count as live_count_routes, models, projects, review, triage)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_app: FastAPI):
|
|
config.ensure_dirs()
|
|
db.migrate()
|
|
from backend import datasets as dataset_store
|
|
adopted = dataset_store.adopt_legacy_tree()
|
|
if adopted:
|
|
print(f"[startup] moved {adopted} legacy dataset tree(s) under datasets/<id>/")
|
|
from backend import projects as project_store
|
|
project_store.ensure_seed_project()
|
|
interrupted = jobs.recover()
|
|
if interrupted:
|
|
print(f"[startup] closed {interrupted} job(s) interrupted by the last restart")
|
|
yield
|
|
|
|
|
|
app = FastAPI(title="Dataset Enrichment", lifespan=lifespan)
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=config.CORS_ORIGINS,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
app.include_router(projects.router)
|
|
app.include_router(batches.router)
|
|
app.include_router(review.router)
|
|
app.include_router(models.router)
|
|
app.include_router(job_routes.router)
|
|
app.include_router(triage.router)
|
|
app.include_router(dataset_routes.router)
|
|
app.include_router(live_count_routes.router)
|
|
app.include_router(counting_bench_routes.router)
|
|
|
|
|
|
@app.get("/api/health")
|
|
def health() -> dict:
|
|
import torch
|
|
from backend import hardware
|
|
|
|
free_vram = hardware.free_vram_gb()
|
|
needed = hardware.SAM3_RESIDENT_GB + hardware.SAM3_HEADROOM_GB
|
|
|
|
return {
|
|
"device": "cuda" if torch.cuda.is_available() else "cpu",
|
|
"gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
|
|
"vram_free_gb": free_vram,
|
|
"sam3_ready": free_vram >= needed or _engine_loaded(),
|
|
"ffmpeg": shutil.which("ffmpeg") is not None,
|
|
"ffprobe": shutil.which("ffprobe") is not None,
|
|
"hf_token": bool(os.environ.get("HUGGING_FACE_HUB_TOKEN")),
|
|
"db": db.healthy(),
|
|
"data_dir": config.DATA_DIR,
|
|
"video_root": config.VIDEO_ROOT,
|
|
"model_loaded": _engine_loaded(),
|
|
}
|
|
|
|
|
|
|
|
def _engine_loaded() -> bool:
|
|
from backend.sam3_engine import engine_is_loaded
|
|
|
|
return engine_is_loaded()
|