File size: 2,548 Bytes
f5b860e d78af0a f5b860e | 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 | import base64
import logging
from io import BytesIO
from typing import List
import clip
import numpy as np
import torch
from PIL import Image
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
logging.basicConfig(
format='%(asctime)s.%(msecs)03d %(levelname)-8s %(message)s',
level=logging.DEBUG,
datefmt='%Y-%m-%d %H:%M:%S'
)
model, preprocess = clip.load("models/ViT-B-32.pt")
model.cpu().eval()
app = FastAPI()
class Texts(BaseModel):
texts_source: List[str]
texts_target: List[str]
class Images(BaseModel):
images: List[str]
texts: List[str]
class CLIPResponse(BaseModel):
similarity: List[List[float]]
@app.post("/clip_image_to_text", response_model=CLIPResponse, tags=["CLIP"])
def clip_image_to_text(data: Images) -> CLIPResponse:
preprocessed_images = []
for image in data.images:
if 'base64,' not in image:
raise HTTPException(422, "Image must be in base64")
image = BytesIO(base64.b64decode(image.split('base64,')[-1]))
image = Image.open(image)
preprocessed_images.append(preprocess(image))
image_input = torch.tensor(np.stack(preprocessed_images)).cpu()
text_tokens = clip.tokenize(["This is the " + desc for desc in data.texts]).cpu()
with torch.no_grad():
image_features = model.encode_image(image_input).float()
text_features = model.encode_text(text_tokens).float()
image_features /= image_features.norm(dim=-1, keepdim=True)
text_features /= text_features.norm(dim=-1, keepdim=True)
similarity = text_features.cpu().numpy() @ image_features.cpu().numpy().T
logging.debug(f"Similarity: {similarity}")
return CLIPResponse(similarity=similarity.tolist())
@app.post("/clip_text_to_text", response_model=CLIPResponse, tags=["CLIP"])
def clip_text_to_text(data: Texts) -> CLIPResponse:
text_input = clip.tokenize([f"This is {text}" for text in data.texts_source]).cpu()
text_output = clip.tokenize(data.texts_target).cpu()
with torch.no_grad():
input_features = model.encode_text(text_input).float()
output_features = model.encode_text(text_output).float()
input_features /= input_features.norm(dim=-1, keepdim=True)
output_features /= output_features.norm(dim=-1, keepdim=True)
similarity = output_features.cpu().numpy() @ input_features.cpu().numpy().T
logging.debug(f"Similarity: {similarity}")
return CLIPResponse(similarity=similarity.tolist())
@app.get("/ping", tags=["TEST"])
def ping():
return "pong"
|