utils/visualization.py

import matplotlib.pyplot as plt
from utils.presets import EMOTION_PRESETS


def generate_scatter(raw, cinematic, emotion, drama):
    fig, ax = plt.subplots(figsize=(6, 6))

    # Plot target anchors faintly
    for name, preset in EMOTION_PRESETS.items():
        t = preset["target"]
        ax.scatter(t["V"], t["A"], alpha=0.2, s=100, color="gray")

    # Raw point
    ax.scatter(raw["V"], raw["A"], s=200, color="blue", label="Raw")

    # Cinematic point
    ax.scatter(cinematic["V"], cinematic["A"], s=200, color="red", label="Cinematic")

    # Arrow scaled by drama
    ax.arrow(
        raw["V"],
        raw["A"],
        cinematic["V"] - raw["V"],
        cinematic["A"] - raw["A"],
        head_width=0.02,
        length_includes_head=True,
        color="black"
    )

    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_xlabel("Valence")
    ax.set_ylabel("Arousal")
    ax.set_title(f"{emotion} | Drama: {round(drama,2)}")
    ax.legend()

    plt.tight_layout()
    return fig