utils/color_model.py

import numpy as np
from huggingface_hub import hf_hub_download
import tflite_runtime.interpreter as tflite


# Download TFLite model file
model_path = hf_hub_download(
    repo_id="danielritchie/vibe-color-model",
    filename="vibe_model.tflite"
)

# Load interpreter
interpreter = tflite.Interpreter(model_path=model_path)
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()


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

    interpreter.set_tensor(input_details[0]["index"], input_data)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]["index"])

    r, g, b, e, i = output_data[0]

    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>
    """