File size: 918 Bytes
f218553
 
 
 
 
 
 
 
 
 
4ec9249
 
 
 
f218553
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from fastapi import FastAPI, Request
from pydantic import BaseModel
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import torch
import requests
from io import BytesIO

app = FastAPI()

# use TensorFlow format -- Akash 22 april 25
model = CLIPModel.from_pretrained("laion/CLIP-ViT-B-32-laion2B-s34B-b79K")
processor = CLIPProcessor.from_pretrained("laion/CLIP-ViT-B-32-laion2B-s34B-b79K")


class ImageInput(BaseModel):
    image_url: str

@app.post("/vector")
def vectorize_image(data: ImageInput):
    try:
        response = requests.get(data.image_url)
        image = Image.open(BytesIO(response.content)).convert("RGB")

        inputs = processor(images=image, return_tensors="pt")
        with torch.no_grad():
            features = model.get_image_features(**inputs)
        
        return {"vector": features[0].tolist()}
    except Exception as e:
        return {"error": str(e)}