File size: 1,001 Bytes
9191c48 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
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
|