Output expression as 3D plot

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import matplotlib.pyplot as plt
 import numpy as np
 import spaces
 import torch
@@ -147,6 +148,10 @@ def process_func(x: np.ndarray, sampling_rate: int) -> dict:
         y = y.detach().cpu().numpy()
         results.append(y[0])
 
+    # Plot A/D/V values
+    plot_expression(results[1][0], results[1][1], results[1][2])
+    expression_file = "expression.png"
+    plt.savefig(expression_file)
     return (
         f"{round(100 * results[0][0])} years",  # age
         {
@@ -154,11 +159,12 @@ def process_func(x: np.ndarray, sampling_rate: int) -> dict:
             "male": results[0][2],
             "child": results[0][3],
         },
-        {
-            "arousal": results[1][0],
-            "dominance": results[1][1],
-            "valence": results[1][2],
-        }
+        expression_file,
+        # {
+        #     "arousal": results[1][0],
+        #     "dominance": results[1][1],
+        #     "valence": results[1][2],
+        # }
     )
 
 
@@ -181,6 +187,52 @@ def recognize(input_file):
     return process_func(signal, target_rate)
 
 
+def plot_expression(arousal, dominance, valence):
+    r"""3D pixel plot of arousal, dominance, valence."""
+    # Voxels per dimension
+    voxels = 7
+    # Create voxel grid
+    x, y, z = np.indices((voxels + 1, voxels + 1, voxels + 1))
+    voxel = (
+        (x == round(arousal * voxels))
+        & (y == round(dominance * voxels))
+        & (z == round(valence * voxels))
+    )
+    colors = np.empty(voxel.shape, dtype=object)
+    colors[voxel] = "#fcb06c"
+    ax = plt.figure().add_subplot(projection='3d')
+    ax.voxels(voxel, facecolors=colors, edgecolor='k')
+    ax.set_aspect("equal")
+    ax.set_xlim([0, voxels])
+    ax.set_ylim([0, voxels])
+    ax.set_zlim([0, voxels])
+    ax.set_xlabel("arousal", fontsize="large", labelpad=0)
+    ax.set_ylabel("dominance", fontsize="large", labelpad=0)
+    ax.set_zlabel("valence", fontsize="large", labelpad=0)
+    ax.set_xticks(
+        list(range(voxels + 1)),
+        labels=["low", None, None, None, None, None, None, "high"],
+        rotation=45,
+        rotation_mode="anchor",
+        verticalalignment="bottom",
+    )
+    ax.set_yticks(
+        list(range(voxels + 1)),
+        labels=["low", None, None, None, None, None, None, "high"],
+        rotation=-25,
+        rotation_mode="anchor",
+        verticalalignment="bottom",
+    )
+    ax.set_zticks(
+        list(range(voxels + 1)),
+        labels=["low", None, None, None, None, None, None, "high"],
+        rotation=25,
+        rotation_mode="default",
+        verticalalignment="bottom",
+    )
+
+
+
 description = (
     "Recognize "
     f"[age](https://huggingface.co/{age_gender_model_name}), "
@@ -204,7 +256,8 @@ with gr.Blocks() as demo:
             with gr.Column():
                 output_age = gr.Textbox(label="Age")
                 output_gender = gr.Label(label="Gender")
-                output_expression = gr.Label(label="Expression")
+                # output_expression = gr.Label(label="Expression")
+                output_expression = gr.Image(label="Expression")
 
         outputs = [output_age, output_gender, output_expression]
         submit_btn.click(recognize, input, outputs)

requirements.txt

@@ -1,4 +1,5 @@
 audiofile
 audresample
+matplotlib
 torch
 transformers