Update app.py

app.py

@@ -578,13 +578,16 @@ def create_avatar():
   
 return img
     
-    # Adjust alpha channel of heatmap
-data = np.array(heatmap_img)
-data[:, :, 3] = data[:, :, 3] * 0.5  # Reduce opacity to 50%
-heatmap_img = Image.fromarray(data)
+  # Adjust alpha channel of heatmap
+    data = np.array(heatmap_img)
+    if data.shape[2] == 3:  # If RGB, add an alpha channel
+        data = np.concatenate([data, np.full((data.shape[0], data.shape[1], 1), 255, dtype=np.uint8)], axis=2)
+    data[:, :, 3] = data[:, :, 3] * 0.5  # Reduce opacity to 50%
+    heatmap_img = Image.fromarray(data)
     
-combined_img = Image.alpha_composite(avatar_img.convert('RGBA'), heatmap_img)
-return combined_img
+    combined_img = Image.alpha_composite(avatar_img.convert('RGBA'), heatmap_img)
+    return combined_img
+
 
 # Create and display avatar with heatmap
 st.subheader("Avatar with Sensation Heatmap")