Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import pyvista as pv
+from pyvista import examples
 import numpy as np
 import librosa
 import requests
@@ -8,11 +9,20 @@ 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):
@@ -210,49 +220,95 @@ def generate_image(sentiment_prediction, transcribed_text):
         # Return a fallback image
         return Image.new('RGB', (1024, 512), color='white')
 
-# Function to create a proper texture-mapped sphere visualization using PyVista
+# 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
 def create_texture_and_sphere_preview(image):
     try:
-        # Save image to temporary file
-        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
-            image.save(tmp.name)
-            texture_file = tmp.name
+        # Convert PIL image to numpy array
+        img_array = np.array(image)
+        height, width = img_array.shape[0], img_array.shape[1]
         
-        # Create a sphere with PyVista
-        sphere = pv.Sphere(radius=1, theta_resolution=100, phi_resolution=50)
+        # 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
+        )
         
-        # Load and apply the texture
-        texture = pv.read_texture(texture_file)
+        # Add the equirectangular image to the first subplot
+        fig.add_trace(go.Image(z=img_array), row=1, col=1)
         
-        # 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'
+        # 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)
         
-        # Capture the image
-        img_array = plotter.screenshot(transparent_background=False)
-        plotter.close()
+        # Convert spherical coordinates to Cartesian coordinates
+        x = np.sin(v) * np.cos(u)
+        y = np.sin(v) * np.sin(u)
+        z = np.cos(v)
         
-        # Convert to PIL Image
-        return Image.fromarray(img_array)
+        # 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)
         
-    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))
+        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
         
-        # 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')
+        # Convert colors to Plotly format (normalized to [0,1])
+        surface_colors = texture_colors.astype(float) / 255.0
         
-        return composite
+        # 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)
+        
+        # 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)
+        
+        return fig
+        
+    except Exception as e:
+        print("Error creating texture and sphere preview:", e)
+        return go.Figure()
 
 # Function to get predictions
 def get_predictions(audio_input):
@@ -271,10 +327,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 texture mapped sphere
-    sphere_visualization = create_texture_and_sphere_preview(image)
+    # Create visualization with both texture and sphere
+    preview_fig = create_texture_and_sphere_preview(image)
     
-    return emotion_prediction, transcribed_text, f"Sentiment: {sentiment} (Polarity: {polarity:.2f})", sphere_visualization, music_path
+    return emotion_prediction, transcribed_text, f"Sentiment: {sentiment} (Polarity: {polarity:.2f})", image, music_path, preview_fig
 
 # Create the Gradio interface
 interface = gr.Interface(
@@ -284,11 +340,12 @@ 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="Texture Mapped Sphere Visualization"),
-        gr.Audio(label="Generated Music", type="filepath")
+        gr.Image(type='pil', label="Generated Equirectangular Image"),
+        gr.Audio(label="Generated Music", type="filepath"),
+        gr.Plot(label="Texture and Sphere Preview")
     ],
     title="Affective Virtual Environments",
-    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."
+    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."
 )
 
 interface.launch()