danielritchie's picture
init
9191c48
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"]