| --- |
| license: mit |
| language: |
| - en |
| pipeline_tag: zero-shot-image-classification |
| tags: |
| - vision |
| - simple |
| - small |
| --- |
| |
| # tinyvvision 🧠✨ |
|
|
| **tinyvvision** is a compact, synthetic curriculum-trained vision-language model designed to demonstrate real zero-shot capability in a minimal setup. Despite its small size (~630k parameters), it aligns images and captions effectively by learning shared visual-language embeddings. |
|
|
| ## What tinyvvision can do: |
| - Match simple geometric shapes (circles, stars, hearts, triangles, etc.) and descriptive captions (e.g., "a red circle", "a yellow star"). |
| - Perform genuine zero-shot generalization, meaning it can correctly match captions to shapes and colors it has never explicitly encountered during training. |
|
|
| ## Model Details: |
| - **Type**: Contrastive embedding (CLIP-style, zero-shot) |
| - **Parameters**: ~630,000 (tiny!) |
| - **Training data**: Fully synthetic—randomly generated shapes, letters, numbers, and symbols paired with descriptive text captions. |
| - **Architecture**: |
| - **Image Encoder**: Simple CNN |
| - **Text Encoder**: Small embedding layer + bidirectional GRU |
| - **Embedding Dim**: 128-dimensional shared embedding space |
|
|
| ## Examples of Zero-Shot Matching: |
| - **Seen during training**: "a red circle" → correctly matches the drawn red circle. |
| - **Never seen**: "a teal lightning bolt" → correctly matched a hand-drawn lightning bolt shape, despite never having seen one during training. |
|
|
| ## Limitations: |
| - tinyvvision is designed as a demonstration of zero-shot embedding and generalization on synthetic data. It is not trained on real-world data or complex scenarios. While robust within its domain (simple geometric shapes and clear captions), results may vary significantly on more complicated or out-of-domain inputs. |
|
|
| ## 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! |
| ```python |
| 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!** ✨ |