File size: 1,004 Bytes
59830d4 | 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 | from pathlib import Path
import pandas as pd
from PIL import Image
from sentence_transformers import SentenceTransformer
from utils.utils import SRC_PATH
image_paths = sorted((SRC_PATH / "data" / "images").glob("*.jpg"))
if not image_paths:
raise FileNotFoundError(f"No JPG images found in {SRC_PATH / 'data'}")
labels = [
"top (tshirt, shirt, polo...)",
]
model = SentenceTransformer(
"sentence-transformers/clip-ViT-B-32",
device="cpu",
)
images = [Image.open(path).convert("RGB") for path in image_paths]
image_embeddings = model.encode(
images,
batch_size=16,
normalize_embeddings=True,
convert_to_numpy=True,
show_progress_bar=True,
)
text_embeddings = model.encode(
prompts,
normalize_embeddings=True,
convert_to_numpy=True,
)
similarities = image_embeddings @ text_embeddings.T
similarity_df = pd.DataFrame(
similarities,
index=[path.name for path in image_paths],
columns=prompts,
)
print(similarity_df.round(3).to_string()) |