Clip

import torchimport clipfrom PIL import Image
# Load modeldevice = "cuda" if torch.cuda.is_available() else "cpu"model, preprocess = clip.load("ViT-B/32", device=device)
# Load imageimage = preprocess(Image.open("photo.jpg")).unsqueeze(0).to(device)
# Define possible labelstext = clip.tokenize(["a dog", "a cat", "a bird", "a car"]).to(device)
# Compute similaritywith torch.no_grad():    image_features = model.encode_image(image)    text_features = model.encode_text(text)
    # Cosine similarity    logits_per_image, logits_per_text = model(image, text)    probs = logits_per_image.softmax(dim=-1).cpu().numpy()
# Print resultslabels = ["a dog", "a cat", "a bird", "a car"]for label, prob in zip(labels, probs[0]):    print(f"{label}: {prob:.2%}")

# Models (sorted by size)models = [    "RN50",           # ResNet-50    "RN101",          # ResNet-101    "ViT-B/32",       # Vision Transformer (recommended)    "ViT-B/16",       # Better quality, slower    "ViT-L/14",       # Best quality, slowest]
model, preprocess = clip.load("ViT-B/32")

# Compute embeddingsimage_features = model.encode_image(image)text_features = model.encode_text(text)
# Normalizeimage_features /= image_features.norm(dim=-1, keepdim=True)text_features /= text_features.norm(dim=-1, keepdim=True)
# Cosine similaritysimilarity = (image_features @ text_features.T).item()print(f"Similarity: {similarity:.4f}")

# Index imagesimage_paths = ["img1.jpg", "img2.jpg", "img3.jpg"]image_embeddings = []
for img_path in image_paths:    image = preprocess(Image.open(img_path)).unsqueeze(0).to(device)    with torch.no_grad():        embedding = model.encode_image(image)        embedding /= embedding.norm(dim=-1, keepdim=True)    image_embeddings.append(embedding)
image_embeddings = torch.cat(image_embeddings)
# Search with text queryquery = "a sunset over the ocean"text_input = clip.tokenize([query]).to(device)with torch.no_grad():    text_embedding = model.encode_text(text_input)    text_embedding /= text_embedding.norm(dim=-1, keepdim=True)
# Find most similar imagessimilarities = (text_embedding @ image_embeddings.T).squeeze(0)top_k = similarities.topk(3)
for idx, score in zip(top_k.indices, top_k.values):    print(f"{image_paths[idx]}: {score:.3f}")

# Define categoriescategories = [    "safe for work",    "not safe for work",    "violent content",    "graphic content"]
text = clip.tokenize(categories).to(device)
# Check imagewith torch.no_grad():    logits_per_image, _ = model(image, text)    probs = logits_per_image.softmax(dim=-1)
# Get classificationmax_idx = probs.argmax().item()max_prob = probs[0, max_idx].item()
print(f"Category: {categories[max_idx]} ({max_prob:.2%})")

# Process multiple imagesimages = [preprocess(Image.open(f"img{i}.jpg")) for i in range(10)]images = torch.stack(images).to(device)
with torch.no_grad():    image_features = model.encode_image(images)    image_features /= image_features.norm(dim=-1, keepdim=True)
# Batch texttexts = ["a dog", "a cat", "a bird"]text_tokens = clip.tokenize(texts).to(device)
with torch.no_grad():    text_features = model.encode_text(text_tokens)    text_features /= text_features.norm(dim=-1, keepdim=True)
# Similarity matrix (10 images × 3 texts)similarities = image_features @ text_features.Tprint(similarities.shape)  # (10, 3)

# Store CLIP embeddings in Chroma/FAISSimport chromadb
client = chromadb.Client()collection = client.create_collection("image_embeddings")
# Add image embeddingsfor img_path, embedding in zip(image_paths, image_embeddings):    collection.add(        embeddings=[embedding.cpu().numpy().tolist()],        metadatas=[{"path": img_path}],        ids=[img_path]    )
# Query with textquery = "a sunset"text_embedding = model.encode_text(clip.tokenize([query]))results = collection.query(    query_embeddings=[text_embedding.cpu().numpy().tolist()],    n_results=5)