Spaces:
Paused
Paused
| """ | |
| Patch-MFFT Tampering Forensics API | |
| ================================== | |
| FastAPI service providing: | |
| - 3-Class Forensics: Authentic vs Full-AI vs Spliced/Inpainted Tampering | |
| - Real-time Pixel-Level Tampering Localization Masks (Base64 PNG) | |
| - Radial Fourier Spectral Band Energy Decomposition | |
| """ | |
| import io | |
| import os | |
| import hashlib | |
| from collections import OrderedDict | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from typing import Dict, List, Optional | |
| from contextlib import asynccontextmanager | |
| from PIL import Image | |
| from fastapi import FastAPI, UploadFile, File, HTTPException, status | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import JSONResponse | |
| from .schemas import HealthResponse, ModelsResponse, PredictionResponse, ModelInfo | |
| from .model_server import TamperingModelServer, predict_ensemble | |
| MODEL_INFO = { | |
| "tiny": {"params": "563K", "description": "Fastest - edge/mobile tampering detection profile (~4 ms)"}, | |
| "base": {"params": "3.28M", "description": "Standard Profile - balanced accuracy and localized inpainting"}, | |
| "large": {"params": "13.94M", "description": "Deep Profile - recommended for vintage/retro AI & degraded media"}, | |
| "ensemble": {"params": "17.22M", "description": "SOTA Production Ensemble - Conservative Fusion of Base + Large with Heatmap Max-Gating"}, | |
| } | |
| PREDICTION_CACHE: OrderedDict = OrderedDict() | |
| MAX_CACHE_SIZE = 256 | |
| model_registry: Dict[str, TamperingModelServer] = {} | |
| DEFAULT_VARIANT = os.environ.get("PATCH_MFFT_VARIANT", "base") | |
| IMG_SIZE = int(os.environ.get("PATCH_MFFT_IMG_SIZE", "384")) | |
| async def lifespan(app: FastAPI): | |
| print("[startup] Initializing Patch-MFFT model registry...") | |
| ckpt_dir_env = os.environ.get("PATCH_MFFT_CHECKPOINT_DIR") | |
| if ckpt_dir_env and Path(ckpt_dir_env).exists(): | |
| ckpt_dir = Path(ckpt_dir_env) | |
| elif (Path(__file__).parent.parent / "checkpoints").exists(): | |
| ckpt_dir = Path(__file__).parent.parent / "checkpoints" | |
| else: | |
| ckpt_dir = Path(__file__).parent.parent / "model" / "checkpoints" | |
| for variant in ["tiny", "base", "large"]: | |
| ckpt_path = ckpt_dir / f"{variant}.pt" | |
| try: | |
| model_registry[variant] = TamperingModelServer( | |
| str(ckpt_path) if ckpt_path.exists() else None, | |
| variant=variant, | |
| img_size=IMG_SIZE, | |
| ) | |
| except Exception as e: | |
| print(f"[startup] Failed to initialize variant '{variant}': {e}") | |
| print(f"[startup] Loaded variants: {list(model_registry.keys())}") | |
| yield | |
| model_registry.clear() | |
| app = FastAPI( | |
| title="Patch-MFFT Tampering Forensics API", | |
| description="Multi-Frequency Transformer for Localized Inpainting, Two-Person Splicing, and Composite Blackmail Detection", | |
| version="1.0.0", | |
| lifespan=lifespan, | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| async def health(): | |
| return HealthResponse( | |
| status="healthy", | |
| version="1.0.0", | |
| models_loaded=list(model_registry.keys()), | |
| timestamp=datetime.now(timezone.utc).isoformat(), | |
| ) | |
| async def list_models(): | |
| default_model = "ensemble" if ("base" in model_registry and "large" in model_registry) else ( | |
| DEFAULT_VARIANT if DEFAULT_VARIANT in model_registry else ( | |
| list(model_registry.keys())[0] if model_registry else "base" | |
| ) | |
| ) | |
| return ModelsResponse( | |
| default=default_model, | |
| models=[ | |
| ModelInfo( | |
| id=v, | |
| loaded=(v in model_registry) or (v == "ensemble" and "base" in model_registry and "large" in model_registry), | |
| params=MODEL_INFO[v]["params"], | |
| description=MODEL_INFO[v]["description"], | |
| ) | |
| for v in MODEL_INFO | |
| ], | |
| ) | |
| async def predict( | |
| file: UploadFile = File(...), | |
| model: str = "ensemble", | |
| ): | |
| variant = (model or "ensemble").lower() | |
| valid_variants = list(model_registry.keys()) + (["ensemble"] if "base" in model_registry and "large" in model_registry else []) | |
| if variant not in valid_variants: | |
| raise HTTPException( | |
| status_code=404, | |
| detail=f"Model variant '{variant}' not found. Available: {valid_variants}", | |
| ) | |
| if not file.content_type or not file.content_type.startswith("image/"): | |
| raise HTTPException(status_code=400, detail="Uploaded file must be a valid image.") | |
| content = await file.read() | |
| if len(content) > 25 * 1024 * 1024: | |
| raise HTTPException(status_code=400, detail="Image size exceeds 25 MB limit.") | |
| sha256_hash = hashlib.sha256(content).hexdigest() | |
| cache_key = f"{sha256_hash}_{variant}" | |
| # Instant In-Memory LRU Cache Check (< 1 ms latency) | |
| if cache_key in PREDICTION_CACHE: | |
| cached_result = dict(PREDICTION_CACHE[cache_key]) | |
| cached_result["cached"] = True | |
| PREDICTION_CACHE.move_to_end(cache_key) | |
| return PredictionResponse(**cached_result) | |
| try: | |
| image = Image.open(io.BytesIO(content)).convert("RGB") | |
| except Exception: | |
| raise HTTPException(status_code=400, detail="Corrupted image file.") | |
| if variant == "ensemble": | |
| result = predict_ensemble( | |
| model_registry["base"], | |
| model_registry["large"], | |
| image, | |
| filename=file.filename, | |
| sha256_hash=sha256_hash, | |
| ) | |
| else: | |
| server = model_registry[variant] | |
| result = server.predict(image, filename=file.filename, sha256_hash=sha256_hash) | |
| # Save to high-speed LRU memory cache | |
| PREDICTION_CACHE[cache_key] = result | |
| if len(PREDICTION_CACHE) > MAX_CACHE_SIZE: | |
| PREDICTION_CACHE.popitem(last=False) | |
| return PredictionResponse(**result) | |