Update app.py

app.py

@@ -471,12 +471,15 @@ def create_heatmap(sensation_map, sensation_type):
     plt.title(f'{["Pain", "Pleasure", "Pressure", "Temperature", "Texture", "EM Field", "Tickle", "Itch", "Quantum", "Neural", "Proprioception", "Synesthesia"][sensation_type]} Sensation Map')
     plt.axis('off')
     
+    # Instead of displaying, save to a buffer
     buf = io.BytesIO()
     plt.savefig(buf, format='png')
     buf.seek(0)
-    plt.close()
+    plt.close()  # Close the figure to free up memory
     
-    return Image.open(buf)
+    # Create an image from the buffer
+    heatmap_img = Image.open(buf)
+    return heatmap_img
 
 def generate_ai_response(keypoints, sensation_map):
     num_keypoints = len(keypoints)
@@ -650,9 +653,12 @@ def create_avatar():
         draw.line([start, end], fill=(0, 255, 255, 50), width=1)
     
     return img
-def create_avatar_with_heatmap(sensation_map, sensation_type):
-    # Create the base avatar
-    avatar_img = create_avatar()
+def create_avatar_with_heatmap(show_heatmap=True):
+    # Load avatar image
+    avatar_img = Image.open("avatar.png").resize((AVATAR_WIDTH, AVATAR_HEIGHT))
+    
+    if not show_heatmap:
+        return avatar_img  # Return the avatar image without heatmap
     
     # Create a heatmap
     heatmap_img = create_heatmap(sensation_map, sensation_type)
@@ -671,9 +677,10 @@ def create_avatar_with_heatmap(sensation_map, sensation_type):
     combined_img = Image.alpha_composite(avatar_img.convert('RGBA'), heatmap_img.convert('RGBA'))
     return combined_img
 
-# Streamlit app
-st.title("NeuraSense AI: Advanced Humanoid Techno-Sensory Simulation")
-
+# Create and display avatar with optional heatmap
+st.subheader("Avatar with Optional Sensation Heatmap")
+avatar_with_heatmap = create_avatar_with_heatmap(show_heatmap)
+st.image(avatar_with_heatmap, use_column_width=True)
 
 
 
@@ -718,6 +725,8 @@ with col2:
     
     # Toggle synesthesia
     use_synesthesia = st.checkbox("Enable Synesthesia", value=False)
+    # Add this with your other UI elements
+    show_heatmap = st.checkbox("Show Sensation Heatmap", value=True)
     
     if canvas_result.json_data is not None:
         objects = canvas_result.json_data["objects"]