| import numpy as np |
| from datasets import load_dataset |
| from sklearn.metrics.pairwise import cosine_similarity |
|
|
| palette_dataset = load_dataset("danielritchie/cinematic-mood-palette")["train"] |
|
|
| palette_vectors = [] |
| palette_names = [] |
|
|
| for row in palette_dataset: |
| palette_vectors.append([row["V"], row["A"], row["D"], row["Cx"], row["Co"]]) |
| palette_names.append(row.get("name", "unknown")) |
|
|
| palette_vectors = np.array(palette_vectors) |
|
|
|
|
| def nearest_palette_vector(raw_vad): |
| raw_vec = np.array([[raw_vad[k] for k in ["V","A","D","Cx","Co"]]]) |
| sims = cosine_similarity(raw_vec, palette_vectors)[0] |
| idx = np.argmax(sims) |
| anchor = palette_vectors[idx] |
|
|
| return { |
| "vector": { |
| "V": anchor[0], |
| "A": anchor[1], |
| "D": anchor[2], |
| "Cx": anchor[3], |
| "Co": anchor[4], |
| }, |
| "name": palette_names[idx] |
| } |
|
|
|
|
| def amplify_with_palette(raw, drama): |
| anchor_data = nearest_palette_vector(raw) |
| anchor = anchor_data["vector"] |
|
|
| amplified = { |
| k: float(raw[k] + drama * (anchor[k] - raw[k])) |
| for k in raw |
| } |
|
|
| return amplified, anchor_data["name"] |
|
|
