Update app.py

app.py

@@ -1,32 +1,39 @@
 import streamlit as st
 import numpy as np
-from typing import List, Tuple
-from transformers import pipeline
 import matplotlib.pyplot as plt
 from matplotlib.backends.backend_agg import FigureCanvasAgg
 from streamlit_drawable_canvas import st_canvas
 import time
 from PIL import Image
 import io
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Constants
-WIDTH, HEIGHT = 600, 600
-ELASTICITY = 0.3
-DAMPING = 0.7
+WIDTH, HEIGHT = 800, 400
+AVATAR_WIDTH, AVATAR_HEIGHT = 300, 400
 
-# Create sensation map
+# Set up DialoGPT model
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
+    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
+    return tokenizer, model
+
+tokenizer, model = load_model()
+
+# Create sensation map for the avatar
 def create_sensation_map(width, height):
     sensation_map = np.zeros((height, width, 3))  # RGB channels for pain, pleasure, and neutral
     for y in range(height):
         for x in range(width):
             # Base sensation
-            base = np.sin(x/30) * np.cos(y/30) * 0.5 + np.random.normal(0, 0.1)
+            base = np.sin(x/15) * np.cos(y/15) * 0.5 + np.random.normal(0, 0.1)
             
             # Pain regions (red channel)
-            pain = np.exp(-((x-150)**2 + (y-150)**2) / 5000) + np.exp(-((x-450)**2 + (y-450)**2) / 5000)
+            pain = np.exp(-((x-75)**2 + (y-100)**2) / 2000) + np.exp(-((x-225)**2 + (y-300)**2) / 2000)
             
             # Pleasure regions (green channel)
-            pleasure = np.exp(-((x-300)**2 + (y-300)**2) / 5000) + np.exp(-((x-150)**2 + (y-450)**2) / 5000)
+            pleasure = np.exp(-((x-150)**2 + (y-200)**2) / 2000) + np.exp(-((x-75)**2 + (y-300)**2) / 2000)
             
             # Neutral sensation (blue channel)
             neutral = 1 - (pain + pleasure)
@@ -35,51 +42,66 @@ def create_sensation_map(width, height):
     
     return sensation_map
 
-sensation_map = create_sensation_map(WIDTH, HEIGHT)
-
-# Set up the Hugging Face pipeline
-@st.cache_resource
-def load_model():
-    return pipeline('text-generation', model='gpt2')
-
-text_generator = load_model()
+avatar_sensation_map = create_sensation_map(AVATAR_WIDTH, AVATAR_HEIGHT)
 
 # Streamlit app
-st.title("Advanced Artificial Touch Simulation")
-
-# Create a Streamlit container for the touch simulation
-touch_container = st.container()
+st.title("Advanced Humanoid Touch Simulation")
+
+# Create two columns
+col1, col2 = st.columns(2)
+
+# Avatar column
+with col1:
+    st.subheader("Humanoid Avatar")
+    avatar_fig, avatar_ax = plt.subplots(figsize=(4, 6))
+    avatar_ax.imshow(avatar_sensation_map)
+    avatar_ax.axis('off')
+    st.pyplot(avatar_fig)
+
+# Touch interface column
+with col2:
+    st.subheader("Touch Interface")
+    touch_fig, touch_ax = plt.subplots(figsize=(4, 6))
+    touch_ax.add_patch(plt.Rectangle((0, 0), AVATAR_WIDTH, AVATAR_HEIGHT, fill=False))
+    touch_ax.set_xlim(0, AVATAR_WIDTH)
+    touch_ax.set_ylim(0, AVATAR_HEIGHT)
+    touch_ax.axis('off')
+    
+    # Convert matplotlib figure to Image
+    canvas = FigureCanvasAgg(touch_fig)
+    canvas.draw()
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    img = Image.open(buf)
+
+    # Use streamlit-drawable-canvas for interaction
+    canvas_result = st_canvas(
+        fill_color="rgba(255, 165, 0, 0.3)",
+        stroke_width=3,
+        stroke_color="#e00",
+        background_color="#eee",
+        background_image=img,
+        update_streamlit=True,
+        height=AVATAR_HEIGHT,
+        width=AVATAR_WIDTH,
+        drawing_mode="point",
+        point_display_radius=5,
+        key="canvas",
+    )
 
 def calculate_sensation(x, y, pressure, duration):
-    # Get sensation from the map
-    sensation = sensation_map[int(y), int(x)]
-    
-    # Modify sensation based on pressure and duration
+    sensation = avatar_sensation_map[int(y), int(x)]
     modified_sensation = sensation * pressure * (1 + np.log(duration + 1))
-    
     return modified_sensation
 
-def on_touch(x, y, pressure, duration):
-    sensation = calculate_sensation(x, y, pressure, duration)
-    pain = sensation[0]
-    pleasure = sensation[1]
-    neutral = sensation[2]
-    
-    # Generate a description of the touch
-    st.write(f"Touch at ({x:.2f}, {y:.2f}) with pressure {pressure:.2f} for {duration:.2f} seconds")
-    st.write(f"Pain: {pain:.2f}, Pleasure: {pleasure:.2f}, Neutral: {neutral:.2f}")
+def generate_description(x, y, pressure, duration, pain, pleasure, neutral):
+    prompt = f"Human: Describe the sensation when touched at ({x:.1f}, {y:.1f}) with pressure {pressure:.2f} for {duration:.2f} seconds. Pain: {pain:.2f}, Pleasure: {pleasure:.2f}, Neutral: {neutral:.2f}.\nAvatar:"
     
-    # Determine the dominant sensation
-    if pain > pleasure and pain > neutral:
-        dominant = "pain"
-    elif pleasure > pain and pleasure > neutral:
-        dominant = "pleasure"
-    else:
-        dominant = "neutral"
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    output = model.generate(input_ids, max_length=150, num_return_sequences=1, no_repeat_ngram_size=2, top_k=50, top_p=0.95, temperature=0.7)
     
-    prompt = f"The user touched the screen at ({x:.2f}, {y:.2f}) with a pressure of {pressure:.2f} for {duration:.2f} seconds, resulting in a {dominant} sensation. Pain: {pain:.2f}, Pleasure: {pleasure:.2f}, Neutral: {neutral:.2f}. Describe the experience:"
-    text = text_generator(prompt, max_length=100, num_return_sequences=1, do_sample=True, top_k=50, top_p=0.95, num_beams=1)[0]['generated_text']
-    st.write(text)
+    return tokenizer.decode(output[0], skip_special_tokens=True).split("Avatar: ")[-1].strip()
 
 # Initialize session state
 if 'touch_start_time' not in st.session_state:
@@ -87,34 +109,6 @@ if 'touch_start_time' not in st.session_state:
 if 'last_touch_position' not in st.session_state:
     st.session_state.last_touch_position = None
 
-# Main app logic
-fig, ax = plt.subplots(figsize=(6, 6))
-ax.imshow(sensation_map)
-ax.axis('off')
-
-# Convert matplotlib figure to Image
-canvas = FigureCanvasAgg(fig)
-canvas.draw()
-buf = io.BytesIO()
-plt.savefig(buf, format='png')
-buf.seek(0)
-img = Image.open(buf)
-
-# Use streamlit-drawable-canvas for interaction
-canvas_result = st_canvas(
-    fill_color="rgba(255, 165, 0, 0.3)",
-    stroke_width=3,
-    stroke_color="#e00",
-    background_color="#eee",
-    background_image=img,
-    update_streamlit=True,
-    height=HEIGHT,
-    width=WIDTH,
-    drawing_mode="point",
-    point_display_radius=0,
-    key="canvas",
-)
-
 # Handle touch events
 if canvas_result.json_data is not None:
     objects = canvas_result.json_data["objects"]
@@ -136,12 +130,21 @@ if canvas_result.json_data is not None:
                 pressure = 1.0
 
             duration = time.time() - st.session_state.touch_start_time
-            on_touch(current_position[0], current_position[1], pressure, duration)
+            x, y = current_position
+            sensation = calculate_sensation(x, y, pressure, duration)
+            pain, pleasure, neutral = sensation
+
+            description = generate_description(x, y, pressure, duration, pain, pleasure, neutral)
+            
+            st.write(f"Touch at ({x:.1f}, {y:.1f}) with pressure {pressure:.2f} for {duration:.2f} seconds")
+            st.write(f"Pain: {pain:.2f}, Pleasure: {pleasure:.2f}, Neutral: {neutral:.2f}")
+            st.write("Avatar's response:")
+            st.write(description)
             
             st.session_state.last_touch_position = current_position
     else:
         st.session_state.touch_start_time = None
         st.session_state.last_touch_position = None
 
-st.write("Click and drag on the image to simulate touch. The color represents different sensations.")
-st.write("Red areas are pain regions, green areas are pleasure regions, and blue areas are neutral.")
+st.write("Click and drag on the touch interface to simulate touching the avatar.")
+st.write("The avatar's sensation map shows pain (red), pleasure (green), and neutral (blue) areas.")