import torch
import torch.nn.functional as F
from io import BytesIO
from urllib.request import Request, urlopen
from PIL import Image
from visme import VisME
model = VisME("path/to/this/repo")
model = model.cuda().eval()
instruction = "Represent the face with the following text."
text = "Retrieve all images with the same cartoon character."
samples = {
"SpiderMan_comic_E616": "https://static.wikia.nocookie.net/spiderman/images/a/ad/Peter_Parker_%28Earth-616%29_017.png/revision/latest?cb=20210807043502",
"SpiderMan_promo_E199999": "https://static.wikia.nocookie.net/marveldatabase/images/2/28/Peter_Parker_%28Earth-199999%29_from_Spider-Man_No_Way_Home_promotional_art_002.jpg/revision/latest/scale-to-width-down/1000?cb=20230730084204",
"Toxin": "https://static.wikia.nocookie.net/superheroes/images/5/53/Toxin.jpg/revision/latest?cb=20240813171932",
}
def load_image(url: str) -> Image.Image:
req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
with urlopen(req, timeout=30) as resp:
return Image.open(BytesIO(resp.read())).convert("RGB")
names = list(samples.keys())
batch = [
{"image": load_image(samples[name]), "text": text, "instruction": instruction}
for name in names
]
with torch.no_grad():
embeddings = model.encode_input(batch)
for i in range(len(names)):
for j in range(i + 1, len(names)):
sim = F.cosine_similarity(embeddings[i : i + 1], embeddings[j : j + 1]).item()
print(f"cosine_similarity({names[i]}, {names[j]}) = {sim:.4f}")