Spaces:
Sleeping
Sleeping
File size: 3,027 Bytes
84ee555 1205230 84ee555 1205230 84ee555 1205230 84ee555 1205230 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import io
import uvicorn
from fastapi import FastAPI, UploadFile, status, Form
from contextlib import asynccontextmanager
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, RedirectResponse
from pydantic import BaseModel
from config import settings
from PIL import Image
from faiss_db import FAISS_DB
from vlm_model import VisionLanguageModel
import torch
import gc
import base64
vectordb = vectordb = FAISS_DB(
faiss_index_file_path=settings.FAISS_INDEX_FILE_PATH,
keyframes_groups_json_path=settings.KEYFRAMES_GROUPS_JSON_PATH,
)
vlm_model = VisionLanguageModel(
model_name=settings.MODEL_NAME,
pretrained=settings.MODEL_CHECKPOINT,
device=(
"cuda" if settings.DEVICE == "cuda" and torch.cuda.is_available() else "cpu"
),
)
@asynccontextmanager
async def lifespan(app: FastAPI):
global vectordb
global vlm_model
yield
del vectordb
del vlm_model
torch.cuda.empty_cache()
gc.collect()
torch.cuda.ipc_collect()
app = FastAPI(title="CLIP Retrieval API", lifespan=lifespan)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_headers=["*"],
allow_credentials=True,
allow_methods=["*"],
)
@app.get("/", include_in_schema=False)
async def root() -> None:
return RedirectResponse("/docs")
@app.get("/health", status_code=status.HTTP_200_OK, tags=["health"])
async def perform_healthcheck() -> None:
return JSONResponse(content={"message": "success"})
class Item(BaseModel):
query: str | UploadFile
queryType: str
topk: int
@app.post("/retrieval")
async def retrieval(
query: str | UploadFile = Form(...),
queryType: str = Form("text"), # text || image
topk: int = Form(...),
) -> JSONResponse:
try:
if queryType == "text":
input_query = query
elif queryType == "image":
image_bytes = await query.read()
input_query = Image.open(io.BytesIO(image_bytes)).convert("RGB")
else:
return JSONResponse(
status_code=status.HTTP_400_BAD_REQUEST,
content={
"message": "Invalid query type. Provide either a text query or an image."
},
)
query_vector = vlm_model.get_embedding(input_query)
except Exception as e:
return JSONResponse(
status_code=status.HTTP_400_BAD_REQUEST,
content={"message": f"Invalid query: {str(e)}"},
)
try:
search_results = vectordb.kNN_search(query_vector, topk)
except Exception as e:
return JSONResponse(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
content={"message": f"Search error: {str(e)}"},
)
return JSONResponse(
status_code=status.HTTP_200_OK,
content={"message": "success", "details": search_results},
)
# Start API
if __name__ == "__main__":
uvicorn.run("main:app", host=settings.HOST, port=settings.PORT, reload=True) |