Update app.py

app.py

@@ -16,6 +16,8 @@ import tempfile
 import base64
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
 
 # Load the emotion prediction model
 def load_emotion_model(model_path):
@@ -52,7 +54,7 @@ processor, music_model, device = load_musicgen_model()
 # Function to transcribe audio
 def transcribe(wav_filepath):
     try:
-        segments, _ = model2.transcribe(wav_filepath, beam_size=5)
+        segments, _ = model2.transscribe(wav_filepath, beam_size=5)
         return "".join([segment.text for segment in segments])
     except Exception as e:
         print("Error transcribing audio:", e)
@@ -186,66 +188,68 @@ def generate_image(emotion_prediction, transcribed_text):
         # Return a fallback image
         return Image.new('RGB', (1024, 512), color='white')
 
-# Function to create a visualization with both the equirectangular image and a 3D sphere
-def create_texture_and_sphere_preview(image):
+# Function to create a sphere with the equirectangular texture using matplotlib
+def create_sphere_with_texture(image):
     try:
-        # Convert PIL image to numpy array for display
+        # Convert PIL image to numpy array
         img_array = np.array(image)
         
-        # Create a subplot with the equirectangular image and a 3D sphere
-        fig = make_subplots(
-            rows=1, cols=2,
-            subplot_titles=("Equirectangular Texture", "3D Sphere Preview"),
-            specs=[[{"type": "image"}, {"type": "scatter3d"}]],
-            horizontal_spacing=0.1
-        )
+        # Create a figure with two subplots
+        fig = plt.figure(figsize=(12, 6))
         
-        # Add the equirectangular image to the first subplot
-        fig.add_trace(go.Image(z=img_array), row=1, col=1)
+        # First subplot: equirectangular image
+        ax1 = fig.add_subplot(121)
+        ax1.imshow(img_array)
+        ax1.set_title('Equirectangular Texture')
+        ax1.axis('off')
         
-        # Create a 3D sphere for the second subplot
-        # Since we can't directly apply the texture, we'll create a colored sphere
-        u = np.linspace(0, 2 * np.pi, 50)
-        v = np.linspace(0, np.pi, 25)
-        u, v = np.meshgrid(u, v)
+        # Second subplot: 3D sphere with texture
+        ax2 = fig.add_subplot(122, projection='3d')
         
-        x = np.sin(v) * np.cos(u)
-        y = np.sin(v) * np.sin(u)
-        z = np.cos(v)
+        # Define a grid matching the map size, subsample along with pixels
+        theta = np.linspace(0, np.pi, img_array.shape[0])
+        phi = np.linspace(0, 2*np.pi, img_array.shape[1])
         
-        # Create a color pattern based on the sphere coordinates
-        colorscale = [(0, 'red'), (0.5, 'green'), (1, 'blue')]
+        # Keep a reasonable number of points for performance
+        count = 100
+        theta_inds = np.linspace(0, img_array.shape[0] - 1, count).round().astype(int)
+        phi_inds = np.linspace(0, img_array.shape[1] - 1, count*2).round().astype(int)
+        theta = theta[theta_inds]
+        phi = phi[phi_inds]
+        img_sampled = img_array[np.ix_(theta_inds, phi_inds)]
         
-        fig.add_trace(go.Surface(
-            x=x, y=y, z=z,
-            surfacecolor=z,  # Use z-coordinate for color
-            colorscale='Viridis',
-            showscale=False,
-            opacity=0.8
-        ), row=1, col=2)
+        # Create meshgrid
+        theta, phi = np.meshgrid(theta, phi)
+        R = 1
         
-        # Update layout
-        fig.update_layout(
-            height=400,
-            title_text="Equirectangular Texture and 3D Sphere Preview",
-            showlegend=False
-        )
+        # Sphere coordinates
+        x = R * np.sin(theta) * np.cos(phi)
+        y = R * np.sin(theta) * np.sin(phi)
+        z = R * np.cos(theta)
         
-        # Update axes for the image subplot
-        fig.update_xaxes(visible=False, row=1, col=1)
-        fig.update_yaxes(visible=False, row=1, col=1)
+        # Plot the sphere with texture
+        ax2.plot_surface(x, y, z, facecolors=img_sampled/255, rstride=1, cstride=1)
         
-        # Update 3D scene settings
-        fig.update_scenes(
-            aspectmode='data',
-            row=1, col=2
-        )
+        # Make the plot more spherical
+        ax2.set_box_aspect([1, 1, 1])  # Aspect ratio is 1:1:1
+        ax2.set_axis_off()
+        ax2.set_title('3D Sphere with Texture')
+        
+        # Adjust viewing angle
+        ax2.view_init(elev=30, azim=45)
         
-        return fig
+        plt.tight_layout()
         
+        # Save the figure to a temporary file
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
+            plt.savefig(tmp_file.name, dpi=100, bbox_inches='tight')
+            plt.close(fig)
+            return tmp_file.name
+            
     except Exception as e:
-        print("Error creating texture and sphere preview:", e)
-        return go.Figure()
+        print("Error creating sphere with texture:", e)
+        # Return a fallback image path
+        return None
 
 # Function to get predictions
 def get_predictions(audio_input):
@@ -265,10 +269,15 @@ def get_predictions(audio_input):
     # Generate music based on transcription and emotion
     music_path = generate_music(transcribed_text, emotion_prediction)
     
-    # Create visualization with both texture and sphere
-    preview_fig = create_texture_and_sphere_preview(image)
+    # Create sphere with texture visualization
+    sphere_image_path = create_sphere_with_texture(image)
+    
+    # Load the sphere image if it was created successfully
+    sphere_image = None
+    if sphere_image_path:
+        sphere_image = Image.open(sphere_image_path)
     
-    return emotion_prediction, transcribed_text, f"Sentiment: {sentiment} (Polarity: {polarity:.2f})", image, music_path, preview_fig
+    return emotion_prediction, transcribed_text, f"Sentiment: {sentiment} (Polarity: {polarity:.2f})", image, music_path, sphere_image
 
 # Create the Gradio interface
 interface = gr.Interface(
@@ -280,10 +289,10 @@ interface = gr.Interface(
         gr.Label(label="Sentiment Analysis"),
         gr.Image(type='pil', label="Generated Equirectangular Image"),
         gr.Audio(label="Generated Music", type="filepath"),
-        gr.Plot(label="Texture and Sphere Preview")
+        gr.Image(type='pil', label="3D Sphere with Texture")
     ],
     title="Affective Virtual Environments",
-    description="Create an AVE using your voice. Get emotion prediction, transcription, sentiment analysis, a generated equirectangular image, music, and a preview of how it would look as a texture on a sphere."
+    description="Create an AVE using your voice. Get emotion prediction, transcription, sentiment analysis, a generated equirectangular image, music, and a 3D sphere with your texture applied."
 )
 
 interface.launch()