Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import pyvista as pv
 import numpy as np
 import librosa
 import requests
@@ -7,15 +8,11 @@ from PIL import Image
 import os
 from tensorflow.keras.models import load_model
 from faster_whisper import WhisperModel
-import random
 from textblob import TextBlob
 import torch
 import scipy.io.wavfile
 from transformers import AutoProcessor, MusicgenForConditionalGeneration
 import tempfile
-import base64
-import plotly.graph_objects as go
-from plotly.subplots import make_subplots
 
 # Load the emotion prediction model
 def load_emotion_model(model_path):
@@ -213,95 +210,49 @@ def generate_image(sentiment_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
-# Function to create a visualization with both the equirectangular image and a 3D sphere
+# Function to create a proper texture-mapped sphere visualization using PyVista
 def create_texture_and_sphere_preview(image):
     try:
-        # Convert PIL image to numpy array
-        img_array = np.array(image)
-        height, width = img_array.shape[0], img_array.shape[1]
+        # Save image to temporary file
+        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
+            image.save(tmp.name)
+            texture_file = tmp.name
         
-        # Create a subplot with the equirectangular image and a 3D sphere
-        fig = make_subplots(
-            rows=1, cols=2,
-            subplot_titles=("Equirectangular Texture", "3D Sphere with Texture Mapping"),
-            specs=[[{"type": "image"}, {"type": "scatter3d"}]],
-            horizontal_spacing=0.1
-        )
-        
-        # Add the equirectangular image to the first subplot
-        fig.add_trace(go.Image(z=img_array), row=1, col=1)
-        
-        # Create sphere coordinates
-        u_res, v_res = 50, 25
-        u = np.linspace(0, 2 * np.pi, u_res)
-        v = np.linspace(0, np.pi, v_res)
-        u, v = np.meshgrid(u, v)
-        
-        # Convert spherical coordinates to Cartesian coordinates
-        x = np.sin(v) * np.cos(u)
-        y = np.sin(v) * np.sin(u)
-        z = np.cos(v)
+        # Create a sphere with PyVista
+        sphere = pv.Sphere(radius=1, theta_resolution=100, phi_resolution=50)
         
-        # Sample colors from the equirectangular image based on UV coordinates
-        # This approximates texture mapping by sampling the image at the correct UV coordinates
-        texture_colors = np.zeros((v_res, u_res, 3), dtype=np.uint8)
+        # Load and apply the texture
+        texture = pv.read_texture(texture_file)
         
-        for i in range(v_res):
-            for j in range(u_res):
-                # Convert spherical coordinates to image coordinates
-                img_x = int((u[i, j] / (2 * np.pi)) * (width - 1))
-                img_y = int((v[i, j] / np.pi) * (height - 1))
-                
-                # Ensure coordinates are within bounds
-                img_x = max(0, min(img_x, width - 1))
-                img_y = max(0, min(img_y, height - 1))
-                
-                # Get color from image
-                if len(img_array.shape) == 3:  # RGB image
-                    texture_colors[i, j] = img_array[img_y, img_x, :3]
-                else:  # Grayscale image
-                    texture_colors[i, j] = [img_array[img_y, img_x]] * 3
+        # Plot with PyVista (off-screen rendering)
+        plotter = pv.Plotter(off_screen=True, window_size=[800, 400])
+        plotter.add_mesh(sphere, texture=texture)
+        plotter.camera_position = 'xy'
+        plotter.camera.azimuth = 30
+        plotter.camera.elevation = 30
+        plotter.background_color = 'white'
         
-        # Convert colors to Plotly format (normalized to [0,1])
-        surface_colors = texture_colors.astype(float) / 255.0
+        # Capture the image
+        img_array = plotter.screenshot(transparent_background=False)
+        plotter.close()
         
-        # Create surface with sampled colors
-        fig.add_trace(go.Surface(
-            x=x, y=y, z=z,
-            surfacecolor=surface_colors,
-            showscale=False,
-            opacity=1.0,
-            lighting=dict(ambient=0.8, diffuse=0.8, specular=0.1, roughness=0.5),
-            lightposition=dict(x=100, y=100, z=100)
-        ), row=1, col=2)
+        # Convert to PIL Image
+        return Image.fromarray(img_array)
         
-        # Update layout
-        fig.update_layout(
-            height=500,
-            title_text="Equirectangular Texture and 3D Sphere Preview",
-            showlegend=False,
-            scene2=dict(
-                xaxis=dict(visible=False, showticklabels=False),
-                yaxis=dict(visible=False, showticklabels=False),
-                zaxis=dict(visible=False, showticklabels=False),
-                aspectmode='data',
-                camera=dict(
-                    eye=dict(x=1.8, y=1.8, z=1.8)
-                ),
-                bgcolor='rgba(0,0,0,0)'
-            )
-        )
-        
-        # Update axes for the image subplot
-        fig.update_xaxes(visible=False, row=1, col=1)
-        fig.update_yaxes(visible=False, row=1, col=1)
+    except Exception as e:
+        print("Error creating texture and sphere preview with PyVista:", e)
+        # Fallback: create a composite image showing both
+        width, height = image.size
+        composite = Image.new('RGB', (width * 2, height), color='white')
+        composite.paste(image, (0, 0))
         
-        return fig
+        # Add text indicating the sphere visualization
+        from PIL import ImageDraw
+        draw = ImageDraw.Draw(composite)
+        draw.text((width + 20, height//2 - 20), "3D Sphere Preview", fill='black')
+        draw.text((width + 20, height//2), "(Texture mapped sphere)", fill='gray')
         
-    except Exception as e:
-        print("Error creating texture and sphere preview:", e)
-        return go.Figure()
+        return composite
 
 # Function to get predictions
 def get_predictions(audio_input):
@@ -320,10 +271,10 @@ def get_predictions(audio_input):
     # Generate music using ACOUSTIC EMOTION prediction with specific prompt
     music_path = generate_music(transcribed_text, emotion_prediction)
     
-    # Create visualization with both texture and sphere
-    preview_fig = create_texture_and_sphere_preview(image)
+    # Create visualization with texture mapped sphere
+    sphere_visualization = create_texture_and_sphere_preview(image)
     
-    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})", sphere_visualization, music_path
 
 # Create the Gradio interface
 interface = gr.Interface(
@@ -333,12 +284,11 @@ interface = gr.Interface(
         gr.Label(label="Acoustic Emotion Prediction (for music)"),
         gr.Label(label="Transcribed Text"),
         gr.Label(label="Sentiment Analysis (for image)"),
-        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="Texture Mapped Sphere Visualization"),
+        gr.Audio(label="Generated Music", type="filepath")
     ],
     title="Affective Virtual Environments",
-    description="Create an AVE using your voice. Get emotion prediction (for music), transcription, sentiment analysis (for image), 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 (for music), transcription, sentiment analysis (for image), a texture-mapped sphere visualization, and generated music."
 )
 
 interface.launch()