73 lines
2.2 KiB
Python
73 lines
2.2 KiB
Python
import warnings
|
|
from dotenv import load_dotenv
|
|
import torch,glob
|
|
from PIL import Image
|
|
from torchvision import transforms
|
|
import os
|
|
from qdrant_client import AsyncQdrantClient, models
|
|
from db_setup import get_qdrant_client
|
|
from vector_db_router.models import CollectionHandler
|
|
from vector_db_router.serializers import CreateCollectionSerializer
|
|
import torch.nn.functional as F
|
|
load_dotenv()
|
|
|
|
# Optional dependency warnings from DINOv2 internals are non-critical.
|
|
warnings.filterwarnings("ignore", message="xFormers is not available.*", category=UserWarning)
|
|
|
|
# Load model
|
|
model = torch.hub.load(
|
|
'facebookresearch/dinov2',
|
|
'dinov2_vitb14'
|
|
)
|
|
|
|
model.eval()
|
|
|
|
# Device
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
model = model.to(device)
|
|
|
|
# Image preprocessing
|
|
transform = transforms.Compose([
|
|
transforms.Resize((518, 518)), # DINOv2 recommended size
|
|
transforms.ToTensor(),
|
|
])
|
|
|
|
def get_embedding(image_path):
|
|
# Load image
|
|
image = Image.open(image_path).convert("RGB")
|
|
|
|
# Transform
|
|
tensor = transform(image).unsqueeze(0).to(device)
|
|
|
|
# Generate embedding
|
|
with torch.no_grad():
|
|
embedding = model(tensor)
|
|
|
|
# Normalize embedding (important for cosine similarity)
|
|
embedding = F.normalize(embedding, p=2, dim=1)
|
|
|
|
# Return flat list (squeeze batch dim)
|
|
return embedding.squeeze(0).cpu().tolist()
|
|
|
|
def get_vectors(image_input, item):
|
|
try:
|
|
base_dir = os.path.join(os.path.dirname(__file__), "downloaded_images")
|
|
path = image_input
|
|
|
|
# If image_input is not a valid file, try to find one using the item (ASIN)
|
|
if not (path and os.path.isfile(path)):
|
|
# If path is a directory, use it as base_dir
|
|
search_base = path if (path and os.path.isdir(path)) else base_dir
|
|
glob_pattern = os.path.join(search_base, item, "*.jpg")
|
|
jpg_files = glob.glob(glob_pattern)
|
|
if jpg_files:
|
|
path = jpg_files[0]
|
|
else:
|
|
return None
|
|
|
|
# Generate the vector for the identified image file
|
|
emb = get_embedding(path)
|
|
return emb.squeeze().tolist()
|
|
except Exception as e:
|
|
print(f"Error generating vector for {item}: {e}")
|
|
return None |