Update app.py

app.py

@@ -6,7 +6,8 @@ import matplotlib.pyplot as plt
 from PIL import Image, ImageDraw, ImageFont
 import time
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import io
+import seaborn as sns
+from io import BytesIO
 import base64
 from streamlit_drawable_canvas import st_canvas
 
@@ -494,9 +495,26 @@ def create_avatar():
   
    return img
 
+# New function to create heatmap
+def create_heatmap(sensation_map):
+    overall_sensitivity = np.mean(sensation_map, axis=2)
+    fig, ax = plt.subplots(figsize=(10, 15))
+    sns.heatmap(overall_sensitivity, cmap='YlOrRd', alpha=0.7, cbar=False, ax=ax)
+    ax.set_axis_off()
+    buf = BytesIO()
+    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+    buf.seek(0)
+    heatmap_img = Image.open(buf)
+    plt.close(fig)
+    return heatmap_img
 
 
-avatar_image = create_avatar()
+def create_avatar_with_heatmap():
+    avatar_img = create_avatar()
+    heatmap_img = create_heatmap(avatar_sensation_map)
+    heatmap_img = heatmap_img.resize((AVATAR_WIDTH, AVATAR_HEIGHT))
+    combined_img = Image.alpha_composite(avatar_img.convert('RGBA'), heatmap_img.convert('RGBA'))
+    return combined_img
 
 
 
@@ -514,18 +532,17 @@ col1, col2 = st.columns([2, 1])
 with col1:
    st.subheader("Humanoid Avatar Interface")
   
-   # Use st_canvas for touch input
-   canvas_result = st_canvas(
-       fill_color="rgba(0, 255, 255, 0.3)",
-       stroke_width=2,
-       stroke_color="#00FFFF",
-       background_image=avatar_image,
-       height=AVATAR_HEIGHT,
-       width=AVATAR_WIDTH,
-       drawing_mode="point",
-       key="canvas",
-   )
-
+  # Use st_canvas for touch input
+    canvas_result = st_canvas(
+        fill_color="rgba(0, 255, 255, 0.3)",
+        stroke_width=2,
+        stroke_color="#00FFFF",
+        background_image=avatar_with_heatmap,
+        height=AVATAR_HEIGHT,
+        width=AVATAR_WIDTH,
+        drawing_mode="point",
+        key="canvas",
+    )
 
 
 # Touch controls and output
@@ -544,13 +561,13 @@ with col2:
    # Toggle synesthesia
    use_synesthesia = st.checkbox("Enable Synesthesia", value=False)
   
-   if canvas_result.json_data is not None:
-       objects = canvas_result.json_data["objects"]
-       if len(objects) > 0:
-           last_touch = objects[-1]
-           touch_x, touch_y = last_touch["left"], last_touch["top"]
-          
-           sensation = avatar_sensation_map[int(touch_y), int(touch_x)]
+     if canvas_result.json_data is not None:
+        objects = canvas_result.json_data["objects"]
+        if len(objects) > 0:
+            last_touch = objects[-1]
+            touch_x, touch_y = last_touch["left"], last_touch["top"]
+           
+         sensation = avatar_sensation_map[int(touch_y), int(touch_x)]
            (
                pain, pleasure, pressure_sens, temp_sens, texture_sens,
                em_sens, tickle_sens, itch_sens, quantum_sens, neural_sens,