utils/color_model.py

import numpy as np
import tensorflow as tf
from huggingface_hub import hf_hub_download

# Download .h5 model from HF
model_path = hf_hub_download(
    repo_id="danielritchie/vibe-color-model",
    filename="vibe_model.h5"
)

# Load Keras model
model = tf.keras.models.load_model(model_path, compile=False)


def infer_color(vad):
    input_data = np.array([[
        vad["V"],
        vad["A"],
        vad["D"],
        vad["Cx"],
        vad["Co"]
    ]], dtype=np.float32)

    output = model.predict(input_data, verbose=0)[0]

    r, g, b, e, i = output

    return {
        "R": float(r),
        "G": float(g),
        "B": float(b),
        "E": float(e),
        "I": float(i)
    }

def scale_rgb(rgb):
    return {
        "R": int(max(0, min(255, rgb["R"] * 255))),
        "G": int(max(0, min(255, rgb["G"] * 255))),
        "B": int(max(0, min(255, rgb["B"] * 255))),
        "E": rgb["E"],
        "I": rgb["I"]
    }

def render_color(rgb):
    return f"""
    <div style="
        width:100%;
        height:240px;
        border-radius:18px;
        background: rgb({rgb['R']},{rgb['G']},{rgb['B']});
        box-shadow: 0px 6px 32px rgba(0,0,0,0.25);
        transition: all 0.3s ease-in-out;
    "></div>
    """