haripra1112001/clip-cifar100-mobilenet

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import torchvision.transforms as transforms
import numpy as np
from huggingface_hub import hf_hub_download
from PIL import Image

class ImageEncoder(nn.Module):
    def __init__(self, proj_dim=128, device="cpu"):
        super().__init__()
        self.device = device
        base = models.mobilenet_v3_small(
            weights=models.MobileNet_V3_Small_Weights.DEFAULT
        )
        self.backbone = nn.Sequential(*list(base.children())[:-1]).to(device)

        for p in self.backbone.parameters():
            p.requires_grad = False

        total_layers = len(list(self.backbone[0].children()))
        for i, layer in enumerate(self.backbone[0].children()):
            if i >= total_layers - 3:
                for p in layer.parameters():
                    p.requires_grad = True

        self.projection = nn.Sequential(
            nn.Linear(576, 1024), nn.BatchNorm1d(1024), nn.ReLU(inplace=True), nn.Dropout(0.20),
            nn.Linear(1024, 512), nn.BatchNorm1d(512),  nn.ReLU(inplace=True), nn.Dropout(0.15),
            nn.Linear(512,  256), nn.BatchNorm1d(256),  nn.ReLU(inplace=True), nn.Dropout(0.10),
            nn.Linear(256, proj_dim),
        ).to(device)

    def forward(self, x):
        feats = self.backbone(x).flatten(1)   # [B, 576]
        out   = self.projection(feats)         # [B, proj_dim]
        return feats, out                      # use 'out', normalise externally

# Download and load
path = hf_hub_download(repo_id="haripra1112001/clip-cifar100-mobilenet",
                       filename="best_cifar100_projection.pth")
ckpt = torch.load(path, map_location="cpu", weights_only=False)

model = ImageEncoder(proj_dim=128, device="cpu")
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

text_emb    = ckpt["text_embeddings"].numpy()   # (100, 128)
class_words = ckpt["class_words"]               # list of 100 CIFAR-100 class names
text_emb_norm = text_emb / np.linalg.norm(text_emb, axis=1, keepdims=True)

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

def predict(image_path, top_k=5):
    img = preprocess(Image.open(image_path).convert("RGB")).unsqueeze(0)
    with torch.no_grad():
        _, proj = model(img)                            # use second output
    img_emb = F.normalize(proj, p=2, dim=1).numpy()    # (1, 128)
    sims    = (text_emb_norm @ img_emb.T).flatten()    # (100,)
    top_idx = np.argsort(-sims)[:top_k]
    return [(class_words[i], round(float(sims[i]), 3)) for i in top_idx]

print(predict("my_image.jpg"))
# e.g. [('bear', 0.42), ('leopard', 0.31), ('wolf', 0.28), ...]

def create_tta_transforms():
    """8 deterministic TTA transforms: center crop, flips, and multiple scales."""
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
    return [
        # View 1: center crop baseline
        transforms.Compose([transforms.Resize(256), transforms.CenterCrop(224),
                            transforms.ToTensor(), normalize]),
        # View 2: center crop + horizontal flip
        transforms.Compose([transforms.Resize(256), transforms.CenterCrop(224),
                            transforms.RandomHorizontalFlip(p=1.0),
                            transforms.ToTensor(), normalize]),
        # Views 3-5: multiple scales
        transforms.Compose([transforms.Resize(232), transforms.CenterCrop(224),
                            transforms.ToTensor(), normalize]),
        transforms.Compose([transforms.Resize(240), transforms.CenterCrop(224),
                            transforms.ToTensor(), normalize]),
        transforms.Compose([transforms.Resize(248), transforms.CenterCrop(224),
                            transforms.ToTensor(), normalize]),
        # Views 6-8: scales + flip
        transforms.Compose([transforms.Resize(232), transforms.CenterCrop(224),
                            transforms.RandomHorizontalFlip(p=1.0),
                            transforms.ToTensor(), normalize]),
        transforms.Compose([transforms.Resize(240), transforms.CenterCrop(224),
                            transforms.RandomHorizontalFlip(p=1.0),
                            transforms.ToTensor(), normalize]),
        transforms.Compose([transforms.Resize(248), transforms.CenterCrop(224),
                            transforms.RandomHorizontalFlip(p=1.0),
                            transforms.ToTensor(), normalize]),
    ]


def apply_tta_to_image(image_tensor, vision_model, device="cpu"):
    """
    Apply 8-view TTA to a single image tensor and return a normalised embedding.

    Args:
        image_tensor: torch.Tensor (C, H, W) or (H, W, C), range [0,1] or [0,255]
        vision_model: loaded ImageEncoder in eval mode
        device: 'cpu' or 'cuda'

    Returns:
        avg_embedding: np.ndarray of shape (1, proj_dim), L2-normalised
    """
    # Convert tensor → PIL Image
    if len(image_tensor.shape) == 3:
        if image_tensor.shape[0] == 3:          # (C, H, W)
            img_np = image_tensor.cpu().permute(1, 2, 0).numpy()
        else:                                    # (H, W, C)
            img_np = image_tensor.cpu().numpy()
    else:
        raise ValueError(f"Expected 3D tensor, got shape {image_tensor.shape}")

    if img_np.max() > 1.0:
        img_np = img_np / 255.0

    pil_image = Image.fromarray((img_np * 255).astype(np.uint8))

    # Apply all 8 transforms and stack into a batch
    tta_batch = torch.stack([t(pil_image) for t in create_tta_transforms()]).to(device)

    # Get embeddings for all 8 views
    with torch.no_grad():
        model_output = vision_model(tta_batch)
        visual_proj = model_output[1] if isinstance(model_output, tuple) else model_output

    # Average UNNORMALISED embeddings, then normalise once
    avg_embedding = F.normalize(visual_proj.mean(dim=0, keepdim=True), p=2, dim=1)
    return avg_embedding.cpu().numpy()   # (1, proj_dim)


def predict_tta(image_path, top_k=5):
    img_tensor = transforms.ToTensor()(Image.open(image_path).convert("RGB"))
    img_emb = apply_tta_to_image(img_tensor, model, device="cpu")   # (1, 128)
    sims    = (text_emb_norm @ img_emb.T).flatten()                 # (100,)
    top_idx = np.argsort(-sims)[:top_k]
    return [(class_words[i], round(float(sims[i]), 3)) for i in top_idx]

print(predict_tta("my_image.jpg"))
# e.g. [('bear', 0.44), ('leopard', 0.32), ('wolf', 0.27), ...]

@misc{prajapati2026clip,
  title   = {CLIP-Style Visual Encoder for CIFAR-100: Contrastive Alignment
             with Visually-Grounded Skip-Gram Embeddings},
  author  = {Prajapati, Harishkumar Kishorkumar},
  year    = {2026}
}