Clip

import sys
# Dependencies are provisioned by this course edition's pinned runtime profile.

import torch
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import os
np.set_printoptions(precision=2, suppress=True)

device = "cuda" if torch.cuda.is_available() else "cpu"

def preprocess(image):
    pixels = np.asarray(image.convert("RGB").resize((32, 32)), dtype=np.float32) / 255.0
    return torch.tensor(pixels).permute(2, 0, 1)

def course_tokenize(prompts):
    rows = []
    for prompt in prompts:
        lowered = prompt.lower()
        rows.append([float("cat" in lowered), float("penguin" in lowered), float("bear" in lowered), min(len(lowered) / 40.0, 1.0)])
    return torch.tensor(rows, dtype=torch.float32)

class CourseClipFixture(torch.nn.Module):
    def encode_image(self, images):
        means = images.mean(dim=(2, 3))
        luminance = images.mean(dim=(1, 2, 3)).unsqueeze(1)
        return torch.cat((means, luminance), dim=1)

    def encode_text(self, tokens):
        return tokens

    def forward(self, images, tokens):
        image_features = torch.nn.functional.normalize(self.encode_image(images), dim=1)
        text_features = torch.nn.functional.normalize(self.encode_text(tokens), dim=1)
        logits = 10.0 * image_features @ text_features.T
        return logits, logits.T

model = CourseClipFixture().to(device)
clip_tokenize = course_tokenize

from pathlib import Path
from PIL import Image
import numpy as np

clip_root = Path("oxcats")
clip_root.mkdir(parents=True, exist_ok=True)
for index, name in enumerate(("Maine_Coon_1.jpg", "Abyssinian_1.jpg", "Persian_1.jpg")):
    yy, xx = np.mgrid[:96, :96]
    pixels = np.stack(((xx * (2 + index) + 60) % 256, (yy * (3 + index) + 40) % 256, ((xx + yy) * 2 + 25 * index) % 256), axis=-1).astype(np.uint8)
    Image.fromarray(pixels).save(clip_root / name)
print("Prepared three generated cat images")

image = preprocess(Image.open("oxcats/Maine_Coon_1.jpg")).unsqueeze(0).to(device)
text = clip_tokenize(["a penguin", "a bear", "a cat"]).to(device)

with torch.no_grad():
    image_features = model.encode_image(image)
    text_features = model.encode_text(text)

    logits_per_image, logits_per_text = model(image, text)
    probs = logits_per_image.softmax(dim=-1).cpu().numpy()

print("Label probs:", probs)

cats_img = [ Image.open(os.path.join("oxcats",x)) for x in os.listdir("oxcats") ]
cats = torch.cat([ preprocess(i).unsqueeze(0) for i in cats_img ]).to(device)
text = clip_tokenize(["a very fat gray cat"]).to(device)
with torch.no_grad():
    logits_per_image, logits_per_text = model(cats, text)
    res = logits_per_text.softmax(dim=-1).argmax().cpu().numpy()

print("Img Index:", res)

plt.imshow(cats_img[res])