31 lines
997 B
Python
31 lines
997 B
Python
import torch
|
|
import open_clip
|
|
from PIL import Image
|
|
import os
|
|
|
|
_model = None
|
|
_preprocess = None
|
|
_device = None
|
|
|
|
def get_model():
|
|
global _model, _preprocess, _device
|
|
if _model is None:
|
|
_device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
# ViT-H-14 is 1024-dim
|
|
_model, _, _preprocess = open_clip.create_model_and_transforms('ViT-H-14', pretrained='laion2b_s32b_b79k')
|
|
_model = _model.to(_device)
|
|
_model.eval()
|
|
return _model, _preprocess, _device
|
|
|
|
def generate_embedding(image_path):
|
|
model, preprocess, device = get_model()
|
|
try:
|
|
image = preprocess(Image.open(image_path)).unsqueeze(0).to(device)
|
|
with torch.no_grad():
|
|
image_features = model.encode_image(image)
|
|
image_features /= image_features.norm(dim=-1, keepdim=True)
|
|
return image_features.cpu().numpy()[0].tolist()
|
|
except Exception as e:
|
|
print(f"Error generating embedding for {image_path}: {e}")
|
|
return None
|