| import io |
| import json |
| import os |
| import asyncio |
| from contextlib import asynccontextmanager |
| from typing import Optional, List |
|
|
| from fastapi import FastAPI, File, Form, UploadFile, HTTPException |
| from fastapi.middleware.cors import CORSMiddleware |
| from fastapi.staticfiles import StaticFiles |
| from fastapi.responses import FileResponse, JSONResponse |
| from PIL import Image |
| import uvicorn |
|
|
| from vit_explainer import ViTExplanationService |
|
|
| |
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) |
| service = ViTExplanationService(base_dir=BASE_DIR) |
| simulate_lock = asyncio.Lock() |
|
|
| def run_inference_task(*args, **kwargs): |
| return service.run_simulation(*args, **kwargs) |
|
|
|
|
|
|
| @asynccontextmanager |
| async def lifespan(app: FastAPI): |
| |
| print("Preloading ViT models into RAM...") |
| try: |
| service.get_model("vit_base_patch16_224") |
| print("✓ ViT-Base (224) loaded into RAM.") |
| except Exception as e: |
| print(f"Notice: Failed to preload ViT-Base: {e}") |
|
|
| try: |
| service.get_model("vit_small_patch16_224") |
| print("✓ ViT-Small (224) loaded into RAM.") |
| except Exception as e: |
| print(f"Notice: Failed to preload ViT-Small: {e}") |
|
|
| print("All models ready in RAM.") |
| yield |
| print("Shutting down API server...") |
|
|
|
|
| app = FastAPI( |
| title="ViT Explanation Comparison API", |
| description="Interactive evaluation for F2HF, BeyondIntuition, Integrated Gradients, Grad-CAM, and Attention Rollout", |
| version="1.0.0", |
| lifespan=lifespan |
| ) |
|
|
| |
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_credentials=True, |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
|
|
| @app.get("/health") |
| @app.get("/api/health") |
| async def health_check(): |
| """Ultra-lightweight ping endpoint to keep cloud containers awake without running inference.""" |
| return {"status": "ok", "message": "Service is healthy"} |
|
|
|
|
| @app.get("/api/classes") |
| async def get_classes(): |
| """Returns the list of all 1000 ImageNet classes for search and selection.""" |
| classes = service.get_classes_list() |
| return {"classes": classes, "total": len(classes)} |
|
|
|
|
| @app.post("/api/simulate") |
| async def simulate( |
| file: UploadFile = File(...), |
| class_query: str = Form("0"), |
| query_type: str = Form("auto"), |
| model_name: str = Form("vit_base_patch16_224"), |
| methods: Optional[str] = Form(None) |
| ): |
| """ |
| Runs model inference and selected explanation methods on the uploaded image. |
| Uses an asyncio lock so simultaneous requests queue cleanly in memory. |
| """ |
| if not file.content_type.startswith("image/"): |
| raise HTTPException(status_code=400, detail="Uploaded file is not a valid image.") |
|
|
| try: |
| image_bytes = await file.read() |
| pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") |
| except Exception as e: |
| raise HTTPException(status_code=400, detail=f"Cannot decode image file: {str(e)}") |
|
|
| |
| selected_methods: Optional[List[str]] = None |
| if methods: |
| try: |
| |
| if methods.strip().startswith("["): |
| selected_methods = json.loads(methods) |
| else: |
| selected_methods = [m.strip().lower() for m in methods.split(",") if m.strip()] |
| except Exception: |
| selected_methods = None |
|
|
| |
| async with simulate_lock: |
| try: |
| loop = asyncio.get_event_loop() |
| result = await loop.run_in_executor( |
| None, |
| run_inference_task, |
| pil_image, |
| class_query, |
| query_type, |
| model_name, |
| selected_methods |
| ) |
| return JSONResponse(content={"status": "success", "data": result}) |
| except Exception as e: |
| import traceback |
| traceback.print_exc() |
| raise HTTPException(status_code=500, detail=f"Explanation calculation failed: {str(e)}") |
|
|
|
|
| |
| static_dir = os.path.join(BASE_DIR, "static") |
| os.makedirs(static_dir, exist_ok=True) |
|
|
|
|
| @app.get("/example_map.json") |
| async def serve_example_map(): |
| map_file = os.path.join(BASE_DIR, "example_map.json") |
| if os.path.exists(map_file): |
| return FileResponse(map_file) |
| alt_file = os.path.join(BASE_DIR, "example", "example_map.json") |
| if os.path.exists(alt_file): |
| return FileResponse(alt_file) |
| raise HTTPException(status_code=404, detail="example_map.json not found") |
|
|
|
|
| @app.get("/") |
| async def serve_index(): |
| index_file = os.path.join(static_dir, "index.html") |
| if os.path.exists(index_file): |
| return FileResponse(index_file) |
| return JSONResponse( |
| content={"message": "ViT Explanation API is running. Frontend static/index.html not yet created."}, |
| status_code=200 |
| ) |
|
|
|
|
| app.mount("/static", StaticFiles(directory=static_dir), name="static") |
|
|
| example_dir = os.path.join(BASE_DIR, "example") |
| if os.path.exists(example_dir): |
| app.mount("/example", StaticFiles(directory=example_dir), name="example") |
|
|
| if __name__ == "__main__": |
| uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=False) |
|
|
