import random import gradio as gr import pandas as pd import numpy as np from Src.Processing import load_data from Src.Processing import process_data from Src.Inference import load_model import torchvision.models as models from torchvision import transforms from Src.Processing_img import ( get_style_model_and_losses, image_loader, run_style_transfer, save_image, gram_matrix ) import torch import tempfile import time # For simulating delays from PIL import Image, ImageDraw, ImageFont # Ensure ImageDraw and ImageFont are imported import os # For file operations import mne import matplotlib.pyplot as plt import io import pyvista as pv import matplotlib.cm as cm import gradio as gr pv.set_plot_theme("document") # A simple theme pv.set_jupyter_backend('html') # --- Data for demonstration --- # Dummy data for Emotion Distribution Bar Chart # In a real app, this would come from your PSD analysis dummy_emotion_data = pd.DataFrame({ 'Emotion': ['sad', 'dis', 'fear', 'neu', 'joy', 'ten', 'ins'], 'Value': [0.8, 0.6, 0.1, 0.4, 0.7, 0.2, 0.3] }) int_to_emotion = { 0: 'sad', 1: 'dis', 2: 'fear', 3: 'neu', 4: 'joy', 5: 'ten', 6: 'ins' } abr_to_emotion = { 'sad': "sadness", 'dis': "disgust", 'fear': "fear", 'neu': "neutral", 'joy': "joy", 'ten': 'Tenderness', 'ins': "inspiration" } # --- Local Image Paths Setup for Dynamic Loading --- # Define a base directory for all painters' images # In Hugging Face Spaces, this would be a folder like 'Painters/' in your repository PAINTERS_BASE_DIR = "Painters" EMOTION_BASE_DIR = "Emotions" model_path = "models\lstm_emotion_model_state.pth" input_size = 320 hidden_size=50 output_size = 7 num_layers=1 # Define painters and some example "filenames" to create placeholders for painters = ["Pablo Picasso", "Vincent van Gogh", "Salvador Dalí"] Base_Dir = "Datasets" # This dictionary defines what placeholder files to create and their captions. # The actual gallery content will be read from the file system. PAINTER_PLACEHOLDER_DATA = { "Pablo Picasso": [ ("Dora Maar with Cat (1941).png", "Dora Maar with Cat (1941)"), ("The Weeping Woman (1937).png", "The Weeping Woman (1937)"), ("Three Musicians (1921).png", "Three Musicians (1921)"), ], "Vincent van Gogh": [ ("Sunflowers (1888).png", "Sunflowers (1888)"), ("The Starry Night (1889).png", "The Starry Night (1889)"), ("The Potato Eaters (1885).png", "The Potato Eaters (1885)"), ], "Salvador Dalí": [ ("Persistence of Memory (1931).png", "Persistence of Memory (1931)"), ("Swans Reflecting Elephants (1937).png", "Swans Reflecting Elephants (1937)"), ("Sleep (1937).png", "Sleep (1937)"), ], } # --- Define the specific PSD files to choose from --- predefined_psd_files = ["task-emotion_psd_1.npy", "task-emotion_psd_2.npy", "task-emotion_psd_3.npy"] # You can put full paths here if they are actual files # --- Core Functions (Simulated) --- def upload_psd_file(selected_file_name): """ Processes a selected PSD file, performs inference, and prepares emotion distribution data. """ if selected_file_name is None: # If no file is selected, return a dummy plot hidden # Return the dummy DataFrame and an empty DataFrame for the state return gr.BarPlot(dummy_emotion_data, x="Emotion", y="Value", label="Emotion Distribution", visible=False), pd.DataFrame() # --- Load and Process PSD Data --- psd_file_path = os.path.join(Base_Dir, selected_file_name).replace(os.sep, '/') # In a real scenario, you'd handle file existence check and errors for load_data try: global np_data np_data = load_data(psd_file_path) print(f"np data orig {np_data.shape}") except FileNotFoundError: print(f"Error: PSD file not found at {psd_file_path}") # Return a plot with error message or just hide it return gr.BarPlot(dummy_emotion_data, x="Emotion", y="Value", label="Emotion Distribution (Error: File not found)", visible=False), pd.DataFrame() final_data = process_data(np_data) # Ensure data is suitable for LSTM (e.g., (batch, sequence_length, input_size)) # If final_data is (sequence_length, input_size), add a batch dimension torch_data = torch.tensor(final_data, dtype=torch.float32).unsqueeze(0) print(f"Processed data shape for model: {torch_data.shape}") # --- Inference --- # Ensure model_path is correct relative to where app.py is run # If 'models' is at your_project_root, adjust path if needed # Assuming 'models' directory is at 'your_project_root' level absolute_model_path = os.path.join("models", "lstm_emotion_model_state.pth") loaded_model = load_model(absolute_model_path, input_size, hidden_size, output_size, num_layers) loaded_model.eval() # Set model to evaluation mode # Pass the prepared torch_data to the model with torch.no_grad(): # Disable gradient calculation for inference predicted_logits, _ = loaded_model(torch_data) # LSTM returns (output, (h_n, c_n)) # Get the most probable emotion index for each time step in the sequence final_output_indices = torch.argmax(predicted_logits, dim=2) # Shape: (batch_size, sequence_length) # Flatten the sequence to count overall emotion frequencies # If batch size is 1, and sequence is long, this view(-1) works for counting all predictions all_predicted_indices = final_output_indices.view(-1) print(f"All predicted indices (flattened): {all_predicted_indices}") # Count occurrences of each predicted emotion index values_count = torch.bincount(all_predicted_indices, minlength=output_size) # Use minlength to ensure all 7 indices are considered print(f"Raw bincount: {values_count}") # --- Create Emotion Distribution DataFrame --- # Initialize emotions_count with all emotions set to 0 frequency emotions_count = {int_to_emotion[i].strip(): 0 for i in range(output_size)} # Use .strip() to remove trailing space from 'sad ' # Update counts only for emotions that were actually predicted for idx, count in enumerate(values_count): if idx < output_size: # Ensure index is within the expected range emotions_count[int_to_emotion[idx].strip()] = count.item() # Use .strip() here too # Convert dictionary to DataFrame dom_emotion = max(emotions_count, key=emotions_count.get) # Ensure column names match what gr.BarPlot expects: "Emotion" and "Frequency" emotion_data = pd.DataFrame({ "Emotion": list(emotions_count.keys()), "Frequency": list(emotions_count.values()) }) # Optional: Sort DataFrame by emotion name or frequency if desired for consistent plotting emotion_data = emotion_data.sort_values(by="Emotion").reset_index(drop=True) print(f"Final emotion_data DataFrame:\n{emotion_data}") # CORRECTED: Return the DataFrame, NOT a gr.BarPlot object # Return both the DataFrame for the plot and the DataFrame itself for the state return gr.BarPlot( emotion_data, x="Emotion", y="Frequency", label="Emotion Distribution", visible=True, y_title="Frequency" ), emotion_data, gr.Textbox(abr_to_emotion[dom_emotion], visible=True) def update_paintings(painter_name): """ Updates the gallery with paintings specific to the selected painter by dynamically listing files in the painter's directory. """ painter_dir = os.path.join(PAINTERS_BASE_DIR, painter_name).replace(os.sep, '/') print(painter_dir) artist_paintings_for_gallery = [] if os.path.isdir(painter_dir): for filename in sorted(os.listdir(painter_dir)): # Sort for consistent order if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')): file_path = os.path.join(painter_dir, filename).replace(os.sep, '/') print(file_path) # Use filename without extension as title, or create more sophisticated parsing title_with_ext = os.path.splitext(filename)[0] artist_paintings_for_gallery.append((file_path, title_with_ext)) print(f"Loaded paintings for {painter_name}: {artist_paintings_for_gallery}") return artist_paintings_for_gallery # Return the list directly for the gallery def generate_my_art(painter, chosen_painting, dom_emotion): """ Simulates the image generation process. 'chosen_painting_info' will be the single selected item from gr.Gallery.select(), e.g., ['Painters/Pablo Picasso/Dora Maar with Cat (1941).png', 'Dora Maar with Cat (1941)'] We need to extract the path. """ print("generating started") print(f"painter: {painter}") print(f"choosen painting: {chosen_painting}") if not painter or not chosen_painting: # Provide default outputs to ensure Gradio components are updated correctly return "Please select a painter and a painting.", None, None ##style image_path img_style_pth = os.path.join(PAINTERS_BASE_DIR, painter, chosen_painting) print(f"img_stype_path: {img_style_pth}") # Display initial status and disable button # --- Simulate your NST or Diffusion Process Here --- # In a real scenario, this would involve your actual NST code. # It would use `selected_painting_path` as the style image. # A content image would be dynamically generated (e.g., a simple colored canvas or # abstract representation based on the PSD analysis's dominant emotion). time.sleep(3) # Simulate processing time # --- Simulate saving a generated image locally --- # This PIL Image would be the actual result of your NST. # We save it to the 'generated_art' directory. """generated_img_pil = Image.new('RGB', (400, 400), color=(np.random.randint(0,255), np.random.randint(0,255), np.random.randint(0,255))) generated_image_local_path = f"generated_art/generated_output_{int(time.time())}.png" generated_img_pil.save(generated_image_local_path) # For the blended image, let's just return the selected style image path for now. # In a real app, this might be a version of the 'generated_img_pil' with a final blend. blended_image_local_path = selected_painting_path """ ##original image emotion_pth = os.path.join(EMOTION_BASE_DIR, dom_emotion) image_name = list(os.listdir(emotion_pth))[random.randint(0, len(os.listdir(emotion_pth)) -1)] original_image_pth = os.path.join(emotion_pth, image_name) print(f"original img _path: {original_image_pth}") final_message = f"Art generated based on {painter}'s {chosen_painting} style!" ## Neural Style Transfer added here device = torch.device("cuda" if torch.cuda.is_available() else "cpu") imsize = 512 if torch.cuda.is_available() else 256 loader = transforms.Compose([ transforms.Resize((imsize, imsize)), transforms.ToTensor() ]) content_layers = ['conv_4'] style_layers = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5'] cnn = models.vgg19(pretrained=True).features.to(device).eval() cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device) cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device) style_img = image_loader(img_style_pth, loader, device) content_img = image_loader(original_image_pth, loader, device) input_img = content_img.clone() output = run_style_transfer(cnn, cnn_normalization_mean, cnn_normalization_std, content_img, style_img, input_img, content_layers, style_layers, device) save_image(output, "stylized_output.jpg") print("Stylized image saved as 'stylized_output.jpg'") stylized_img_path = 'stylized_output.jpg' # Return final results and re-enable button yield gr.Textbox(final_message), original_image_pth, stylized_img_path def generate_topomap(n_channels, n_time): n_sensors = 64 if n_channels is None or n_time is None: print("they are None") n_channels = 4 n_time = 500 # ---------------------------- # 2. Load standard 10-20 montage # ---------------------------- montage = mne.channels.make_standard_montage('standard_1020') # Filter only the standard 64 EEG electrodes standard_64_chs = [ 'Fp1', 'Fp2', 'F7', 'F3', 'Fz', 'F4', 'F8', 'FC5', 'FC1', 'FC2', 'FC6', 'T7', 'C3', 'Cz', 'C4', 'T8', 'CP5', 'CP1', 'CP2', 'CP6', 'P7', 'P3', 'Pz', 'P4', 'P8', 'POz', 'O1', 'Oz', 'O2', 'Fpz', 'AF7', 'AF3', 'AF4', 'AF8', 'F5', 'F1', 'F2', 'F6', 'FC3', 'FCz', 'FC4', 'C5', 'C1', 'C2', 'C6', 'CP3', 'CPz', 'CP4', 'P5', 'P1', 'P2', 'P6', 'PO3', 'PO4', 'PO7', 'PO8', 'PO9', 'PO10', 'O1', 'O2', 'FT7', 'FT8', 'TP7', 'TP8' ] # exactly 64 channels ch_pos_dict = montage.get_positions()['ch_pos'] ch_pos_dict_filtered = {ch: ch_pos_dict[ch] for ch in standard_64_chs} channel_names = list(ch_pos_dict_filtered.keys()) ch_pos_array = np.array([ch_pos_dict_filtered[ch] for ch in standard_64_chs]) # Nx3 ch_pos_2d = ch_pos_array[:, :2] # For 2D topomap # ---------------------------- # 3. Choose a time index and frequency index # ---------------------------- new_data = np_data.reshape(64, 630, 5) print(f"shape: {new_data.shape}") print(f"n channels: {n_channels}") psd_snapshot = new_data[:, n_time - 1, n_channels - 1] # Normalize PSD for coloring psd_norm = (psd_snapshot - psd_snapshot.min()) / (psd_snapshot.max() - psd_snapshot.min()) # ---------------------------- # 4. Plot 2D topomap using MNE # ---------------------------- print(f"shape psd :{psd_snapshot.shape}") fig, ax = plt.subplots() mne.viz.plot_topomap( psd_snapshot, ch_pos_2d, names=channel_names, show=False, axes=ax ) # Save the generated topomap to a temp file tmpfile = tempfile.NamedTemporaryFile(suffix=".png", delete=False) fig.savefig(tmpfile.name, dpi=150, bbox_inches="tight") plt.close(fig) return tmpfile.name predefined_psd_files = ["task-emotion_psd_1.npy", "task-emotion_psd_2.npy", "task-emotion_psd_3.npy"] # --- Gradio Interface Definition --- with gr.Blocks(css=".gradio-container { max-width: 2000px; margin: auto; }") as demo: # Define the gr.State component here, accessible throughout the Blocks # This will hold the information of the SINGLE selected painting from the gallery # This will hold the DataFrame of the emotion distribution (to be passed to generate_my_art) current_emotion_df_state = gr.State(value=pd.DataFrame()) # Header Section gr.Markdown( """

Brain Emotion Decoder 🧠🎨

Imagine seeing your deepest feelings transform into art. We decode the underlying emotions from your brain activity, generating a personalized artwork that comes to life within an interactive 3D brain model. Discover the art of your inner self.

""" ) with gr.Row(): # Left Column: Input and Emotion Distribution with gr.Column(scale=1): gr.Markdown("

1. Choose a PSD file

") # Radio buttons to select from predefined files psd_file_selection = gr.Radio( choices=predefined_psd_files, label="Select a PSD file for analysis", value=predefined_psd_files[0], # Default selection interactive=True ) # Button to trigger PSD analysis analyze_psd_button = gr.Button("Analyze PSD File", variant="secondary") gr.Markdown("

2. Emotion Distribution

") # Bar plot for emotion distribution emotion_distribution_plot = gr.BarPlot( dummy_emotion_data, x="Emotion", y="Value", label="Emotion Distribution", height=300, x_title="Emotion Type", y_title="Frequency", visible=False # Hidden until analysis is triggered ) dom_emotion = gr.Textbox(label = "dominant emotion", visible=False) # Right Column: Art Museum and Generation with gr.Column(scale=1): gr.Markdown("

Your Art Mesum

") # Kept original heading gr.Markdown("

3. Choose your favourite painter

") painter_dropdown = gr.Dropdown( choices=painters, value="Pablo Picasso", # Default selection label="Select a Painter" ) gr.Markdown("

4. Choose your favourite painting

") # Gallery to display paintings for selection painting_gallery = gr.Gallery( # Correct initial value and visibility value=update_paintings("Pablo Picasso"), # Initial load for Picasso's paintings label="Select a Painting", height=300, columns=3, rows=1, object_fit="contain", preview=True, # Allows clicking to see larger image interactive=True, # Make it selectable elem_id="painting_gallery", visible=True, # Should be visible by default ) # Button to trigger art generation selected_painting_name = gr.Textbox(visible=False) generate_button = gr.Button("Generate My Art", variant="primary", scale=0) # Status message for image generation status_message = gr.Textbox( value="Click 'Generate My Art' to begin.", label="Generation Status", interactive=False, show_label=False, lines=1 ) # Output section on a separate "page" or revealed dynamically gr.Markdown( """

Your Generated Artwork

Generated Image") generated_image_output = gr.Image(label="Generated Image", show_label=False, height=300) gr.Markdown("

Blended Style Image

") blended_image_output = gr.Image(label="Blended Style Image", show_label=False, height=300) with gr.Column(scale=1): gr.Markdown("

Brain Topomap

") channels_slider = gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Channels", interactive=True) timestamp_slider = gr.Slider(minimum=1, maximum=630, value=1, step=1, label="Timestamp", interactive=True) mne_2d_img = gr.Image(visible=True) generate_button.click(generate_topomap, outputs=mne_2d_img) # --- Event Listeners --- analyze_psd_button.click( upload_psd_file, inputs=[psd_file_selection], # Input is the selected radio button value (file name) outputs=[emotion_distribution_plot, current_emotion_df_state, dom_emotion] # CORRECTED: Added current_emotion_df_state to outputs ) # When painter dropdown changes, update the gallery content and reset selected_painting_state painter_dropdown.change( update_paintings, # This updates the gallery inputs=[painter_dropdown], outputs=[painting_gallery] # Only output the gallery content directly ) # IMPORTANT: Use the .select() method of gr.Gallery to capture the specific clicked item. # The 'select' event passes the selected value directly as the argument to the function. # We use a lambda to simply return that selected value and store it in our state. def on_select(evt: gr.SelectData): print("this function started") print(f"Image index: {evt.index}\nImage value: {evt.value['image']['orig_name']}") return evt.value['image']['orig_name'] painting_gallery.select( on_select, # This lambda receives the selected image info (path, title) outputs=[selected_painting_name] # The output updates our gr.State component ) # The generate_button now correctly uses the value from selected_painting_state generate_button.click( generate_my_art, inputs=[painter_dropdown, selected_painting_name, dom_emotion], # Pass painter and the SELECTED painting outputs=[status_message, generated_image_output, blended_image_output] ) ## sliders event listener channels_slider.change(fn=generate_topomap, inputs=[channels_slider, timestamp_slider], outputs=mne_2d_img) timestamp_slider.change(fn=generate_topomap, inputs=[channels_slider, timestamp_slider], outputs=mne_2d_img) # Launch the demo if __name__ == "__main__": # Ensure project_root_dir is defined for this block if you uncomment these lines # project_root_dir = os.path.dirname(os.path.abspath(__file__)) # project_root_dir = os.path.dirname(project_root_dir) # print(f"Loading LSTM model from: {os.path.join(project_root_dir, model_path)}") # _ = load_model(os.path.join(project_root_dir, model_path), input_size, hidden_size, output_size, num_layers) demo.launch()