1from huggingface_hub import hf_hub_download
2import torch, re, numpy as np, math
3from PIL import Image, ImageDraw, ImageFont
4
5repo = "ProCreations/tinyvvision"
6pth = hf_hub_download(repo, "cortexclip-mini.pth")
7device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
8state = torch.load(pth, map_location=device)
9idx2tok = state["vocab"]
10tok2idx = {t:i for i,t in enumerate(idx2tok)}
11def encode_txt(s, maxlen=16):
12 toks = re.findall(r"\w+|[^\w\s]", s.lower())
13 ids = [tok2idx.get(t,0) for t in toks][:maxlen]
14 return ids + [0]*(maxlen-len(ids))
15class TE(torch.nn.Module):
16 def __init__(self):
17 super().__init__()
18 self.emb = torch.nn.Embedding(len(idx2tok), 64)
19 self.gru = torch.nn.GRU(64, 128, num_layers=2, bidirectional=True, batch_first=True)
20 self.out_proj = torch.nn.Linear(256, 128)
21 def forward(self, x):
22 e, _ = self.gru(self.emb(x))
23 return self.out_proj(e[:, -1])
24class IE(torch.nn.Module):
25 def __init__(self):
26 super().__init__()
27 self.conv = torch.nn.Sequential(
28 torch.nn.Conv2d(3,32,5,1,2), torch.nn.ReLU(),
29 torch.nn.Conv2d(32,64,3,1,1), torch.nn.ReLU(),
30 torch.nn.Conv2d(64,128,3,1,1), torch.nn.ReLU(),
31 torch.nn.AdaptiveAvgPool2d((4,4)), torch.nn.Flatten(),
32 torch.nn.Linear(128*4*4,128), torch.nn.ReLU()
33 )
34 def forward(self, x): return self.conv(x)
35te, ie = TE().to(device), IE().to(device)
36te.load_state_dict(state["text_encoder"])
37ie.load_state_dict(state["image_encoder"])
38te.eval(); ie.eval()
39
40# ----- CUSTOMIZE YOUR EXAMPLES HERE -----
41# To try your own image:
42# 1. Replace the 'custom_image()' function with your image drawing/loading code.
43# 2. Replace 'custom_caption' with your own caption for the image.
44def custom_image():
45 # Example: Draw your own "blue hexagon" shape below!
46 img = Image.new("RGB",(64,64),"white")
47 dr = ImageDraw.Draw(img)
48 dr.regular_polygon((32,32,22), n_sides=6, fill="blue")
49 arr = np.array(img).astype(np.float32)/255.0
50 return torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).to(device)
51custom_caption = "a blue hexagon"
52
53# ----- FUN DEMO EXAMPLES -----
54def draw_red_heart():
55 img = Image.new("RGB",(64,64),"white")
56 dr = ImageDraw.Draw(img)
57 dr.polygon([(32,18),(50,34),(32,56),(14,34)], fill="red") # simple heart
58 dr.ellipse((18,12,32,32), fill="red")
59 dr.ellipse((32,12,46,32), fill="red")
60 arr = np.array(img).astype(np.float32)/255.0
61 return torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).to(device)
62def draw_purple_star():
63 img = Image.new("RGB",(64,64),"white")
64 dr = ImageDraw.Draw(img)
65 points = [ (32+20*math.cos(math.radians(a)),32+20*math.sin(math.radians(a))) for a in range(-90, 270, 72) ]
66 for i in range(5):
67 dr.line([points[i], points[(i+2)%5]], fill="purple", width=7)
68 arr = np.array(img).astype(np.float32)/255.0
69 return torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).to(device)
70def draw_orange_pentagon():
71 img = Image.new("RGB",(64,64),"white")
72 dr = ImageDraw.Draw(img)
73 dr.regular_polygon((32,32,22), n_sides=5, fill="orange")
74 arr = np.array(img).astype(np.float32)/255.0
75 return torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).to(device)
76
77demo_imgs = [
78 (custom_image(), custom_caption),
79 (draw_red_heart(), "a red heart"),
80 (draw_purple_star(), "a purple star"),
81 (draw_orange_pentagon(), "an orange pentagon"),
82]
83captions = [c for (_,c) in demo_imgs]
84img_tensors = [im for (im,_) in demo_imgs]
85cap_ids = torch.tensor([encode_txt(c) for c in captions], device=device)
86
87with torch.no_grad():
88 txt_emb = te(cap_ids)
89 for i, (img, caption) in enumerate(zip(img_tensors, captions)):
90 im_emb = ie(img)
91 sim = torch.nn.functional.cosine_similarity(im_emb, txt_emb).cpu().numpy()
92 rank = int(np.argmax(sim))
93 print(f"Input image {i+1}: '{caption}'")
94 print(" Similarity scores:")
95 for j, c in enumerate(captions):
96 print(f" {c}: {sim[j]:.4f}")
97 print(" Best match:", captions[rank], "\n")
