MisbahKhan/R2-Tuning

1import os2import tempfile3from fastapi import FastAPI, UploadFile, File, Form, HTTPException4from fastapi.middleware.cors import CORSMiddleware5import clip6import decord7import nncore8import numpy as np9import torch10import pandas as pd11import torchvision.transforms.functional as F12from decord import VideoReader13from nncore.engine import load_checkpoint14from nncore.nn import build_model15from contextlib import asynccontextmanager16 17# Global variables for model and config18model, cfg = None, None19 20# Lifespan handler to manage startup and shutdown21@asynccontextmanager22async def lifespan(app: FastAPI):23    # Startup: Load the model and config24    global model, cfg25    print("Loading model on startup...")26    model, cfg = init_model(CONFIG, WEIGHT)27    print("Model loaded successfully.")28    yield  # Application runs here29    # Shutdown: Clean up (if needed)30    print("Shutting down...")31 32# Initialize FastAPI app with lifespan33app = FastAPI(title="R2-Tuning API", lifespan=lifespan)34 35# Enable CORS for React app36app.add_middleware(37    CORSMiddleware,38    allow_origins=["http://localhost:3000"],39    allow_credentials=True,40    allow_methods=["*"],41    allow_headers=["*"],42)43 44# Configuration45CONFIG = 'configs/qvhighlights/r2_tuning_qvhighlights.py'46WEIGHT = 'r2_tuning_qvhighlights-ed516355.pth'47 48def convert_time(seconds):49    minutes, seconds = divmod(round(max(seconds, 0)), 60)50    return f'{minutes:02d}:{seconds:02d}'51 52def load_video(video_path, cfg):53    decord.bridge.set_bridge('torch')54    vr = VideoReader(video_path)55    stride = vr.get_avg_fps() / cfg.data.val.fps56    fm_idx = [min(round(i), len(vr) - 1) for i in np.arange(0, len(vr), stride).tolist()]57    video = vr.get_batch(fm_idx).permute(0, 3, 1, 2).float() / 25558    size = 336 if '336px' in cfg.model.arch else 22459    h, w = video.size(-2), video.size(-1)60    s = min(h, w)61    x, y = round((h - s) / 2), round((w - s) / 2)62    video = video[..., x:x + s, y:y + s]63    video = F.resize(video, size=(size, size))64    video = F.normalize(video, (0.481, 0.459, 0.408), (0.269, 0.261, 0.276))65    return video.reshape(video.size(0), -1).unsqueeze(0)66 67def init_model(config, checkpoint):68    cfg = nncore.Config.from_file(config)69    cfg.model.init = True70    model = build_model(cfg.model, dist=False).eval()71    model = load_checkpoint(model, checkpoint, warning=False)72    return model, cfg73 74def process_video(video_path: str, query: str, model, cfg) -> dict:75    if not query:76        raise ValueError("Text query cannot be empty.")77    try:78        video = load_video(video_path, cfg)79    except Exception as e:80        raise ValueError(f"Failed to load video: {str(e)}")81    query = clip.tokenize(query, truncate=True)82    device = next(model.parameters()).device83    data = dict(video=video.to(device), query=query.to(device), fps=[cfg.data.val.fps])84    with torch.inference_mode():85        pred = model(data)86    mr = pred['_out']['boundary'][:5].cpu().tolist()87    mr = [[convert_time(p[0]), convert_time(p[1]), round(p[2], 2)] for p in mr]88    hd = pred['_out']['saliency'].cpu()89    hd = ((hd - hd.min()) / (hd.max() - hd.min()) * 0.9 + 0.05).tolist()90    hd = [{"x": i * 2, "y": y} for i, y in enumerate(hd)]91    return {"moment_retrieval": mr, "highlight_detection": hd}92 93@app.post("/predict")94async def predict(video: UploadFile = File(...), query: str = Form(...)):95    try:96        if not video.content_type.startswith("video/"):97            raise HTTPException(status_code=400, detail="Invalid file type. Please upload a video.")98        if not query.strip():99            raise HTTPException(status_code=400, detail="Text query cannot be empty.")100        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_file:101            temp_file.write(await video.read())102            temp_file_path = temp_file.name103        try:104            result = process_video(temp_file_path, query, model, cfg)105            return result106        finally:107            os.unlink(temp_file_path)108    except ValueError as e:109        raise HTTPException(status_code=400, detail=str(e))110    except Exception as e:111        raise HTTPException(status_code=500, detail=f"Server error: {str(e)}")112 113@app.get("/health")114async def health():115    return {"status": "healthy"}