creek

  1"""
  2Emotions
  3Original file: https://colab.research.google.com/drive/1SgZgWv7uK3kibT3KW1zR5jp5pH46AQE0
  4"""
  5
  6!pip install git+https://github.com/openai/CLIP.git --q
  7!pip install torch torchvision pillow ftfy regex tqdm --q
  8
  9import torch
 10import clip
 11from PIL import Image
 12import numpy as np
 13
 14device = "cuda" if torch.cuda.is_available() else "cpu"
 15model, preprocess = clip.load("ViT-B/32", device=device)
 16
 17print("Device:", device)
 18
 19EMOTIONAL_CLASS = [
 20    "sensual",
 21    "playful",
 22    "edgy",
 23    "bold",
 24    "luxurious",
 25    "cheeky",
 26    "ironic",
 27    "mysterious",
 28]
 29
 30EMOTIONAL_PHRASES = {
 31    "sensual": "a sensual emotional tone, conveying intimacy and softness",
 32    "playful": "a playful and fun emotional mood with cheerful energy",
 33    "edgy": "an edgy emotional feel with intensity and attitude",
 34    "bold": "a bold, fearless emotional expression with strong presence",
 35    "luxurious": "a luxurious emotional mood with elegance and refinement",
 36    "cheeky": "a cheeky, mischievous emotional attitude with confidence",
 37    "ironic": "an ironic, detached emotional vibe with witty contrast",
 38    "mysterious": "a mysterious emotional tone filled with intrigue and ambiguity",
 39}
 40
 41
 42EMOTION_PROMPT_TEMPLATES = [
 43    "a photo that conveys {}",
 44    "a portrait expressing {}",
 45    "a cinematic shot capturing {}",
 46    "an artistic image filled with {}",
 47    "a dramatic photograph embodying {}",
 48    "a close-up shot that reveals {}",
 49    "a high-fashion editorial with {}",
 50    "a lifestyle photo characterized by {}",
 51    "a conceptual art image exploring {}",
 52    "a professional photoshoot showcasing {}",
 53]
 54
 55
 56def build_text_bank(output_path="lighting_text_bank"):
 57    all_prompt_texts = []
 58    all_prompt_labels = []
 59    all_prompt_features = []
 60
 61    class_centroids = []
 62
 63    model.eval()
 64    with torch.no_grad():
 65        for lighting in EMOTIONAL_CLASS:
 66            phrase = EMOTIONAL_PHRASES[lighting]
 67
 68            prompts = [tpl.format(phrase) for tpl in EMOTION_PROMPT_TEMPLATES]
 69            tokens = clip.tokenize(prompts).to(device)
 70
 71            text_features = model.encode_text(tokens)
 72            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
 73
 74            all_prompt_texts.extend(prompts)
 75            all_prompt_labels.extend([lighting] * len(prompts))
 76            all_prompt_features.append(text_features.cpu().numpy())
 77
 78            centroid = text_features.mean(dim=0)
 79            centroid = centroid / centroid.norm()
 80            class_centroids.append(centroid.cpu().numpy())
 81
 82    all_prompt_features = np.concatenate(all_prompt_features, axis=0)
 83    class_centroids = np.stack(class_centroids, axis=0)
 84
 85    np.savez(
 86        output_path,
 87        prompt_texts=np.array(all_prompt_texts),
 88        prompt_labels=np.array(all_prompt_labels),
 89        prompt_features=all_prompt_features,
 90        class_centroids=class_centroids,
 91        class_names=np.array(EMOTIONAL_CLASS),
 92    )
 93    print(f"Saved text bank → {output_path}")
 94
 95
 96build_text_bank()
 97
 98
 99def get_image_embedding(image_path: str) -> torch.Tensor:
100    image = Image.open(image_path).convert("RGB")
101    img_tensor = preprocess(image).unsqueeze(0).to(device)
102
103    model.eval()
104    with torch.no_grad():
105        img_feat = model.encode_image(img_tensor)
106        img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True)
107
108    return img_feat.squeeze(0)
109
110
111DEFAULT_THRESHOLD = 0.22
112
113
114def load_text_bank(path="/content/lighting_text_bank.npz"):
115    bank = np.load(path, allow_pickle=True)
116    centroids = torch.tensor(bank["class_centroids"]).to(device)
117    class_names = bank["class_names"]
118
119    centroids = centroids / centroids.norm(dim=-1, keepdim=True)
120    return centroids, class_names
121
122
123def classify_lighting(image_path, threshold=DEFAULT_THRESHOLD):
124    img_feat = get_image_embedding(image_path).to(device)
125
126    centroids, class_names = load_text_bank()
127
128    img_feat = img_feat.to(centroids.dtype)
129
130    img_feat = img_feat / img_feat.norm()
131
132    sims = centroids @ img_feat
133    best_score, best_idx = torch.max(sims, dim=0)
134
135    best_score = float(best_score.detach().cpu())
136    sims_list = sims.detach().cpu().tolist()
137
138    sims_dict = {str(cls): float(s) for cls, s in zip(class_names, sims_list)}
139
140    if best_score < threshold:
141        predicted = "no_specific_emotion"
142    else:
143        predicted = str(class_names[best_idx])
144
145    return predicted, best_score, sims_dict
146
147
148# Runner
149
150from google.colab import files
151
152uploaded = files.upload()
153
154image_path = list(uploaded.keys())[0]
155
156pred, score, sims = classify_lighting(image_path)
157
158print("Prediction:", pred)
159print("Best score:", score)
160print("\nAll similarities:")
161for k, v in sims.items():
162    print(f"{k:20s}: {v:.3f}")