creek

"""
Emotions
Original file: https://colab.research.google.com/drive/1SgZgWv7uK3kibT3KW1zR5jp5pH46AQE0
"""

!pip install git+https://github.com/openai/CLIP.git --q
!pip install torch torchvision pillow ftfy regex tqdm --q

import torch
import clip
from PIL import Image
import numpy as np

device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

print("Device:", device)

EMOTIONAL_CLASS = [
    "sensual",
    "playful",
    "edgy",
    "bold",
    "luxurious",
    "cheeky",
    "ironic",
    "mysterious",
]

EMOTIONAL_PHRASES = {
    "sensual": "a sensual emotional tone, conveying intimacy and softness",
    "playful": "a playful and fun emotional mood with cheerful energy",
    "edgy": "an edgy emotional feel with intensity and attitude",
    "bold": "a bold, fearless emotional expression with strong presence",
    "luxurious": "a luxurious emotional mood with elegance and refinement",
    "cheeky": "a cheeky, mischievous emotional attitude with confidence",
    "ironic": "an ironic, detached emotional vibe with witty contrast",
    "mysterious": "a mysterious emotional tone filled with intrigue and ambiguity",
}


EMOTION_PROMPT_TEMPLATES = [
    "a photo that conveys {}",
    "a portrait expressing {}",
    "a cinematic shot capturing {}",
    "an artistic image filled with {}",
    "a dramatic photograph embodying {}",
    "a close-up shot that reveals {}",
    "a high-fashion editorial with {}",
    "a lifestyle photo characterized by {}",
    "a conceptual art image exploring {}",
    "a professional photoshoot showcasing {}",
]


def build_text_bank(output_path="lighting_text_bank"):
    all_prompt_texts = []
    all_prompt_labels = []
    all_prompt_features = []

    class_centroids = []

    model.eval()
    with torch.no_grad():
        for lighting in EMOTIONAL_CLASS:
            phrase = EMOTIONAL_PHRASES[lighting]

            prompts = [tpl.format(phrase) for tpl in EMOTION_PROMPT_TEMPLATES]
            tokens = clip.tokenize(prompts).to(device)

            text_features = model.encode_text(tokens)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)

            all_prompt_texts.extend(prompts)
            all_prompt_labels.extend([lighting] * len(prompts))
            all_prompt_features.append(text_features.cpu().numpy())

            centroid = text_features.mean(dim=0)
            centroid = centroid / centroid.norm()
            class_centroids.append(centroid.cpu().numpy())

    all_prompt_features = np.concatenate(all_prompt_features, axis=0)
    class_centroids = np.stack(class_centroids, axis=0)

    np.savez(
        output_path,
        prompt_texts=np.array(all_prompt_texts),
        prompt_labels=np.array(all_prompt_labels),
        prompt_features=all_prompt_features,
        class_centroids=class_centroids,
        class_names=np.array(EMOTIONAL_CLASS),
    )
    print(f"Saved text bank → {output_path}")


build_text_bank()


def get_image_embedding(image_path: str) -> torch.Tensor:
    image = Image.open(image_path).convert("RGB")
    img_tensor = preprocess(image).unsqueeze(0).to(device)

    model.eval()
    with torch.no_grad():
        img_feat = model.encode_image(img_tensor)
        img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True)

    return img_feat.squeeze(0)


DEFAULT_THRESHOLD = 0.22


def load_text_bank(path="/content/lighting_text_bank.npz"):
    bank = np.load(path, allow_pickle=True)
    centroids = torch.tensor(bank["class_centroids"]).to(device)
    class_names = bank["class_names"]

    centroids = centroids / centroids.norm(dim=-1, keepdim=True)
    return centroids, class_names


def classify_lighting(image_path, threshold=DEFAULT_THRESHOLD):
    img_feat = get_image_embedding(image_path).to(device)

    centroids, class_names = load_text_bank()

    img_feat = img_feat.to(centroids.dtype)

    img_feat = img_feat / img_feat.norm()

    sims = centroids @ img_feat
    best_score, best_idx = torch.max(sims, dim=0)

    best_score = float(best_score.detach().cpu())
    sims_list = sims.detach().cpu().tolist()

    sims_dict = {str(cls): float(s) for cls, s in zip(class_names, sims_list)}

    if best_score < threshold:
        predicted = "no_specific_emotion"
    else:
        predicted = str(class_names[best_idx])

    return predicted, best_score, sims_dict


# Runner

from google.colab import files

uploaded = files.upload()

image_path = list(uploaded.keys())[0]

pred, score, sims = classify_lighting(image_path)

print("Prediction:", pred)
print("Best score:", score)
print("\nAll similarities:")
for k, v in sims.items():
    print(f"{k:20s}: {v:.3f}")