Update README.md

README.md

@@ -35,5 +35,101 @@ tags:
 
 ## How to Test tinyvvision:
 Check out the provided inference script to easily test your own shapes and captions. Feel free to challenge tinyvvision with new, unseen combinations to explore its generalization capability!
+'from huggingface_hub import hf_hub_download
+import torch, re, numpy as np, math
+from PIL import Image, ImageDraw, ImageFont
 
+repo = "ProCreations/tinyvvision"
+pth = hf_hub_download(repo, "cortexclip-mini.pth")
+device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
+state = torch.load(pth, map_location=device)
+idx2tok = state["vocab"]
+tok2idx = {t:i for i,t in enumerate(idx2tok)}
+def encode_txt(s, maxlen=16):
+    toks = re.findall(r"\w+|[^\w\s]", s.lower())
+    ids = [tok2idx.get(t,0) for t in toks][:maxlen]
+    return ids + [0]*(maxlen-len(ids))
+class TE(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.emb = torch.nn.Embedding(len(idx2tok), 64)
+        self.gru = torch.nn.GRU(64, 128, num_layers=2, bidirectional=True, batch_first=True)
+        self.out_proj = torch.nn.Linear(256, 128)
+    def forward(self, x):
+        e, _ = self.gru(self.emb(x))
+        return self.out_proj(e[:, -1])
+class IE(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Sequential(
+            torch.nn.Conv2d(3,32,5,1,2), torch.nn.ReLU(),
+            torch.nn.Conv2d(32,64,3,1,1), torch.nn.ReLU(),
+            torch.nn.Conv2d(64,128,3,1,1), torch.nn.ReLU(),
+            torch.nn.AdaptiveAvgPool2d((4,4)), torch.nn.Flatten(),
+            torch.nn.Linear(128*4*4,128), torch.nn.ReLU()
+        )
+    def forward(self, x): return self.conv(x)
+te, ie = TE().to(device), IE().to(device)
+te.load_state_dict(state["text_encoder"])
+ie.load_state_dict(state["image_encoder"])
+te.eval(); ie.eval()
+
+# ----- CUSTOMIZE YOUR EXAMPLES HERE -----
+# To try your own image:
+#   1. Replace the 'custom_image()' function with your image drawing/loading code.
+#   2. Replace 'custom_caption' with your own caption for the image.
+def custom_image():
+    # Example: Draw your own "blue hexagon" shape below!
+    img = Image.new("RGB",(64,64),"white")
+    dr = ImageDraw.Draw(img)
+    dr.regular_polygon((32,32,22), n_sides=6, fill="blue")
+    arr = np.array(img).astype(np.float32)/255.0
+    return torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).to(device)
+custom_caption = "a blue hexagon"
+
+# ----- FUN DEMO EXAMPLES -----
+def draw_red_heart():
+    img = Image.new("RGB",(64,64),"white")
+    dr = ImageDraw.Draw(img)
+    dr.polygon([(32,18),(50,34),(32,56),(14,34)], fill="red")  # simple heart
+    dr.ellipse((18,12,32,32), fill="red")
+    dr.ellipse((32,12,46,32), fill="red")
+    arr = np.array(img).astype(np.float32)/255.0
+    return torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).to(device)
+def draw_purple_star():
+    img = Image.new("RGB",(64,64),"white")
+    dr = ImageDraw.Draw(img)
+    points = [ (32+20*math.cos(math.radians(a)),32+20*math.sin(math.radians(a))) for a in range(-90, 270, 72) ]
+    for i in range(5):
+        dr.line([points[i], points[(i+2)%5]], fill="purple", width=7)
+    arr = np.array(img).astype(np.float32)/255.0
+    return torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).to(device)
+def draw_orange_pentagon():
+    img = Image.new("RGB",(64,64),"white")
+    dr = ImageDraw.Draw(img)
+    dr.regular_polygon((32,32,22), n_sides=5, fill="orange")
+    arr = np.array(img).astype(np.float32)/255.0
+    return torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).to(device)
+
+demo_imgs = [
+    (custom_image(), custom_caption),
+    (draw_red_heart(), "a red heart"),
+    (draw_purple_star(), "a purple star"),
+    (draw_orange_pentagon(), "an orange pentagon"),
+]
+captions = [c for (_,c) in demo_imgs]
+img_tensors = [im for (im,_) in demo_imgs]
+cap_ids  = torch.tensor([encode_txt(c) for c in captions], device=device)
+
+with torch.no_grad():
+    txt_emb = te(cap_ids)
+    for i, (img, caption) in enumerate(zip(img_tensors, captions)):
+        im_emb = ie(img)
+        sim = torch.nn.functional.cosine_similarity(im_emb, txt_emb).cpu().numpy()
+        rank = int(np.argmax(sim))
+        print(f"Input image {i+1}: '{caption}'")
+        print("  Similarity scores:")
+        for j, c in enumerate(captions):
+            print(f"    {c}: {sim[j]:.4f}")
+        print("  Best match:", captions[rank], "\n")'
 ✨ **Enjoy experimenting!** ✨