88 lines
2.7 KiB
Python
88 lines
2.7 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,
|
|
sam3, 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(sam3.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()
|