|
|
""" |
|
|
Image embedding service for generating vector embeddings from images |
|
|
""" |
|
|
|
|
|
import io |
|
|
import os |
|
|
import base64 |
|
|
from typing import List, Tuple |
|
|
|
|
|
import open_clip |
|
|
import torch |
|
|
from fastapi import UploadFile, HTTPException |
|
|
from PIL import Image |
|
|
import torch |
|
|
|
|
|
|
|
|
class ImageEmbeddingModel: |
|
|
"""Class for handling image embedding using CLIP model""" |
|
|
|
|
|
def __init__(self, model_name: str): |
|
|
self.model_name = model_name |
|
|
self.device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
self.model, self.preprocess_train, self.preprocess_val = self._initialize_model() |
|
|
|
|
|
def _initialize_model(self) -> Tuple: |
|
|
"""Initialize the CLIP model for image embeddings""" |
|
|
model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms(self.model_name) |
|
|
tokenizer = open_clip.get_tokenizer(self.model_name) |
|
|
model.to(self.device) |
|
|
model.eval() |
|
|
return model, preprocess_train, preprocess_val |
|
|
|
|
|
def get_embedding_from_pil(self, image: Image.Image) -> List[float]: |
|
|
"""Get embedding from PIL image""" |
|
|
processed_image = self.preprocess_val(image).unsqueeze(0).to(self.device) |
|
|
|
|
|
|
|
|
if self.device == 'cuda': |
|
|
autocast_context = torch.amp.autocast(device_type='cuda') |
|
|
else: |
|
|
|
|
|
autocast_context = torch.amp.autocast(device_type='cpu', dtype=torch.float32) |
|
|
with torch.no_grad(), autocast_context: |
|
|
image_features = self.model.encode_image(processed_image, normalize=True) |
|
|
|
|
|
|
|
|
return image_features.cpu().numpy()[0].tolist() |
|
|
|
|
|
async def get_embedding_from_upload(self, image_file: UploadFile) -> List[float]: |
|
|
"""Get embedding from uploaded image file""" |
|
|
try: |
|
|
contents = await image_file.read() |
|
|
img = Image.open(io.BytesIO(contents)).convert("RGB") |
|
|
return self.get_embedding_from_pil(img) |
|
|
except Exception as e: |
|
|
raise HTTPException(status_code=400, detail=f"Invalid image: {str(e)}") |
|
|
|
|
|
def get_embedding_from_base64(self, base64_data: str) -> List[float]: |
|
|
"""Get embedding from base64 encoded image""" |
|
|
try: |
|
|
|
|
|
if ',' in base64_data: |
|
|
base64_data = base64_data.split(',')[1] |
|
|
|
|
|
image_bytes = base64.b64decode(base64_data) |
|
|
image = Image.open(io.BytesIO(image_bytes)).convert("RGB") |
|
|
|
|
|
return self.get_embedding_from_pil(image) |
|
|
except Exception as e: |
|
|
raise HTTPException(status_code=400, detail=f"Invalid base64 image: {str(e)}") |
|
|
|
|
|
def get_embeddings_from_folder(self, image_folder: str) -> List[List[float]]: |
|
|
"""Get embeddings from all images in a folder""" |
|
|
embeddings = [] |
|
|
|
|
|
if not os.path.exists(image_folder): |
|
|
raise HTTPException(status_code=404, detail=f"Folder not found: {image_folder}") |
|
|
|
|
|
for image_name in os.listdir(image_folder): |
|
|
if image_name.lower().endswith(('.png', '.jpg', '.jpeg')): |
|
|
try: |
|
|
image_path = os.path.join(image_folder, image_name) |
|
|
img = Image.open(image_path).convert("RGB") |
|
|
embeddings.append(self.get_embedding_from_pil(img)) |
|
|
except Exception as e: |
|
|
print(f"Error processing {image_name}: {str(e)}") |
|
|
|
|
|
return embeddings |
|
|
|
