Upload 5 files

README.md

@@ -4,7 +4,7 @@ emoji: 🧠
 colorFrom: blue
 colorTo: pink
 sdk: gradio
-sdk_version: 5.20.1
+sdk_version: 3.36.1
 app_file: app.py
 pinned: false
 ---
@@ -25,7 +25,7 @@ This application implements a VAE model that:
 
 This demo uses the [SreekarB/OSFData](https://huggingface.co/datasets/SreekarB/OSFData) dataset from HuggingFace, which contains:
 
-- Functional connectivity matrices from fMRI data
+- NIfTI files in P01_rs.nii format containing fMRI data
 - Demographic information directly in the dataset:
   - ID: Subject identifier 
   - wab_aq: Aphasia quotient score (severity measure)
@@ -35,19 +35,24 @@ This demo uses the [SreekarB/OSFData](https://huggingface.co/datasets/SreekarB/O
   - gender: Subject gender
   - handedness: Subject handedness (ignored in this analysis)
 
+The application processes the NIfTI files using the Power 264 atlas to create functional connectivity matrices that are then analyzed by the VAE model.
+
 ## How to Use
 
-1. **Data Source**: By default, it uses the HuggingFace dataset. You can change to a local directory if needed.
-2. **Model Parameters**:
-   - Latent Dimensions: Controls the size of the latent space (default: 32)
-   - Number of Epochs: Training iterations (default: 1000)
-   - Batch Size: Training batch size (default: 16)
-
-3. **Run the Analysis**: The model will:
-   - Load and process the data
-   - Train the VAE model
-   - Analyze relationships between latent variables and demographics
-   - Generate visualizations of original, reconstructed, and generated FC matrices
+1. **Configure Parameters**:
+   - **Data Source**: By default, it uses the SreekarB/OSFData HuggingFace dataset
+   - **Latent Dimensions**: Controls the size of the latent space (default: 32)
+   - **Number of Epochs**: Training iterations (default: 200 for demo)
+   - **Batch Size**: Training batch size (default: 16)
+
+2. **Start Training**:
+   - Click the "Start Training" button to begin the analysis
+   - The training progress will be displayed in the Status area
+
+3. **View Results**:
+   - The VAE will learn latent representations of brain connectivity
+   - Results will show correlations between demographic variables and latent brain patterns
+   - The visualization shows original FC, reconstructed FC, and a new FC matrix generated from specific demographic values
 
 ## Outputs
 

app.py

@@ -3,6 +3,7 @@ from main import run_fc_analysis
 import os
 
 def gradio_fc_analysis(data_source, latent_dim, nepochs, bsize, use_hf_dataset):
+    """Run the full VAE analysis pipeline"""
     fig = run_fc_analysis(
         data_dir=data_source,
         demographic_file=None,  # We're now getting demographics directly from the dataset
@@ -12,30 +13,17 @@ def gradio_fc_analysis(data_source, latent_dim, nepochs, bsize, use_hf_dataset):
         save_model=True,
         use_hf_dataset=use_hf_dataset
     )
-    return fig
+    return fig, "Analysis complete! VAE model has been trained and demographic relationships analyzed."
 
 def create_interface():
-    iface = gr.Interface(
-        fn=gradio_fc_analysis,
-        inputs=[
-            gr.Textbox(label="Data Source (HF Dataset ID or Local Directory)", 
-                       value="SreekarB/OSFData"),
-            gr.Slider(minimum=8, maximum=64, step=8, 
-                      label="Latent Dimensions", value=32),
-            gr.Slider(minimum=100, maximum=5000, step=100, 
-                      label="Number of Epochs", value=500),  # Reduced for faster demos
-            gr.Slider(minimum=8, maximum=64, step=8, 
-                      label="Batch Size", value=16),
-            gr.Checkbox(label="Use HuggingFace Dataset", 
-                       value=True),
-        ],
-        outputs="plot",
-        title="Aphasia fMRI to FC Analysis using VAE",
-        description="""
-        Analysis pipeline: fMRI → FC matrices → VAE → Analysis
-        
-        This demo uses the SreekarB/OSFData dataset from HuggingFace by default.
-        The dataset contains the following columns:
+    with gr.Blocks(title="Aphasia fMRI to FC Analysis using VAE") as iface:
+        gr.Markdown("""
+        # Aphasia fMRI to FC Analysis using VAE
+        
+        This demo uses a Variational Autoencoder (VAE) to analyze functional connectivity patterns in the brain and their relationship to demographic variables.
+        
+        ## Dataset Information
+        By default, this uses the SreekarB/OSFData dataset from HuggingFace with the following variables:
         - ID: Subject identifier
         - wab_aq: Aphasia severity score
         - age: Age of the subject
@@ -43,12 +31,67 @@ def create_interface():
         - education: Years of education
         - gender: Subject gender
         - handedness: Subject handedness (ignored in the analysis)
-        """,
-        examples=[
-            ["SreekarB/OSFData", 32, 200, 16, True],  # Fewer epochs for faster demo
-        ],
-        cache_examples=False,
-    )
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                # Configuration parameters
+                data_source = gr.Textbox(
+                    label="Data Source (HF Dataset ID or Local Directory)", 
+                    value="SreekarB/OSFData"
+                )
+                latent_dim = gr.Slider(
+                    minimum=8, maximum=64, step=8, 
+                    label="Latent Dimensions", value=32
+                )
+                nepochs = gr.Slider(
+                    minimum=100, maximum=5000, step=100, 
+                    label="Number of Epochs", value=200  # Reduced for faster demos
+                )
+                bsize = gr.Slider(
+                    minimum=8, maximum=64, step=8, 
+                    label="Batch Size", value=16
+                )
+                use_hf_dataset = gr.Checkbox(
+                    label="Use HuggingFace Dataset", value=True
+                )
+                
+                # Training button
+                train_button = gr.Button("Start Training", variant="primary")
+                status_text = gr.Textbox(label="Status", value="Ready to start training")
+                
+            with gr.Column(scale=2):
+                # Output plot
+                output_plot = gr.Plot(label="Analysis Results")
+        
+        # Link the training button to the analysis function
+        train_button.click(
+            fn=gradio_fc_analysis,
+            inputs=[data_source, latent_dim, nepochs, bsize, use_hf_dataset],
+            outputs=[output_plot, status_text]
+        )
+        
+        # Add examples
+        gr.Examples(
+            examples=[
+                ["SreekarB/OSFData", 32, 200, 16, True],  # Fewer epochs for faster demo
+            ],
+            inputs=[data_source, latent_dim, nepochs, bsize, use_hf_dataset],
+        )
+        
+        # Add explanation of the workflow
+        gr.Markdown("""
+        ## How this works
+        
+        1. **Data Loading**: The system downloads NIfTI files (P01_rs.nii format) from the SreekarB/OSFData dataset
+        2. **Preprocessing**: The fMRI data is processed using the Power 264 atlas and converted to functional connectivity (FC) matrices
+        3. **VAE Training**: A conditional VAE model learns the latent representation of brain connectivity
+        4. **Analysis**: The system analyzes relationships between latent brain connectivity patterns and demographic variables
+        5. **Visualization**: Results are displayed showing original FC, reconstructed FC, generated FC, and demographic correlations
+        
+        Note: This app works with the SreekarB/OSFData dataset that contains NIfTI files and demographic information.
+        """)
+    
     return iface
 
 if __name__ == "__main__":

data_preprocessing.py

@@ -4,13 +4,153 @@ from datasets import load_dataset
 from nilearn import input_data, connectome
 from nilearn.image import load_img
 import nibabel as nib
+import os
 
-def preprocess_fmri_to_fc(dataset_name, atlas_path=None):
-    dataset = load_dataset(dataset_name, split="train")
+def preprocess_fmri_to_fc(dataset_or_niifiles, demo_data=None, demo_types=None):
+    """
+    Process fMRI data to generate functional connectivity matrices
     
-    # Load Power 264 atlas or specified atlas
-    if atlas_path is None:
-        # Use Power 264 coordinates to create spherical ROIs
+    Parameters:
+    - dataset_or_niifiles: Either a dataset name string or a list of NIfTI files
+    - demo_data: Optional demographic data, required if providing NIfTI files
+    - demo_types: Optional demographic data types, required if providing NIfTI files
+    
+    Returns:
+    - X: Array of FC matrices
+    - demo_data: Demographic data
+    - demo_types: Demographic data types
+    """
+    print(f"Preprocessing data with type: {type(dataset_or_niifiles)}")
+    
+    # For SreekarB/OSFData dataset, the data will be loaded from dataset features
+    if isinstance(dataset_or_niifiles, str) and dataset_or_niifiles == "SreekarB/OSFData":
+        print("Loading data from SreekarB/OSFData dataset")
+        dataset = load_dataset(dataset_or_niifiles, split="train")
+        
+        # Prepare demographics data from the dataset
+        if demo_data is None:
+            # Create demo_data from the dataset
+            demo_df = pd.DataFrame({
+                'age': dataset['age'],
+                'gender': dataset['gender'],
+                'mpo': dataset['mpo'],
+                'wab_aq': dataset['wab_aq']
+            })
+            
+            demo_data = [
+                demo_df['age'].values,
+                demo_df['gender'].values,
+                demo_df['mpo'].values,
+                demo_df['wab_aq'].values
+            ]
+            
+            demo_types = ['continuous', 'categorical', 'continuous', 'continuous']
+        
+        # Look for NIfTI files in P01_rs.nii format
+        print("Searching for NIfTI files in dataset columns...")
+        nii_files = []
+        
+        # First check if there are any columns with .nii files
+        for col in dataset.column_names:
+            # Check if column contains file paths
+            first_val = dataset[col][0] if len(dataset) > 0 else None
+            if isinstance(first_val, str) and (first_val.endswith('.nii') or first_val.endswith('.nii.gz')):
+                print(f"Found column '{col}' with NIfTI file paths")
+                
+                # Try to download files from HuggingFace Hub
+                from huggingface_hub import hf_hub_download
+                import tempfile
+                
+                for item in dataset[col]:
+                    try:
+                        # Download the NIfTI file
+                        file_path = hf_hub_download(
+                            repo_id="SreekarB/OSFData",
+                            filename=item,
+                            repo_type="dataset"
+                        )
+                        nii_files.append(file_path)
+                        print(f"Downloaded {item} from HuggingFace Hub")
+                    except Exception as e:
+                        print(f"Error downloading {item}: {e}")
+                        # Try looking for the file locally
+                        local_path = os.path.join(os.getcwd(), item)
+                        if os.path.exists(local_path):
+                            nii_files.append(local_path)
+                            print(f"Found {item} locally")
+                        else:
+                            print(f"Warning: Could not find {item} locally or on HuggingFace")
+                
+                # If we found NIfTI files, process them to FC matrices
+                if nii_files:
+                    print(f"Found {len(nii_files)} NIfTI files, converting to FC matrices")
+                    
+                    # Load Power 264 atlas
+                    from nilearn import datasets
+                    power = datasets.fetch_coords_power_2011()
+                    coords = np.vstack((power.rois['x'], power.rois['y'], power.rois['z'])).T
+                    
+                    masker = input_data.NiftiSpheresMasker(
+                        coords, radius=5,
+                        standardize=True,
+                        memory='nilearn_cache', memory_level=1,
+                        verbose=0,
+                        detrend=True,
+                        low_pass=0.1,
+                        high_pass=0.01,
+                        t_r=2.0  # Adjust TR according to your data
+                    )
+                    
+                    # Process fMRI data and compute FC matrices
+                    fc_matrices = []
+                    for nii_file in nii_files:
+                        try:
+                            fmri_img = load_img(nii_file)
+                            time_series = masker.fit_transform(fmri_img)
+                            
+                            correlation_measure = connectome.ConnectivityMeasure(
+                                kind='correlation',
+                                vectorize=False,
+                                discard_diagonal=False
+                            )
+                            
+                            fc_matrix = correlation_measure.fit_transform([time_series])[0]
+                            
+                            triu_indices = np.triu_indices_from(fc_matrix, k=1)
+                            fc_triu = fc_matrix[triu_indices]
+                            
+                            fc_triu = np.arctanh(fc_triu)  # Fisher z-transform
+                            
+                            fc_matrices.append(fc_triu)
+                            print(f"Processed {nii_file} to FC matrix")
+                        except Exception as e:
+                            print(f"Error processing {nii_file}: {e}")
+                    
+                    if fc_matrices:
+                        X = np.array(fc_matrices)
+                        # Normalize the FC data
+                        X = (X - np.mean(X, axis=0)) / np.std(X, axis=0)
+                        print(f"Created FC matrices with shape {X.shape}")
+                        return X, demo_data, demo_types
+                break  # Stop after finding one column with NIfTI files
+        
+        # If we're here, we couldn't find or process the NIfTI files
+        print("Could not find or process NIfTI files from the dataset.")
+        
+        print("No FC or fMRI data found in the dataset. Please provide FC matrices.")
+        # Return a placeholder with the right demographics but empty FC
+        n_subjects = len(dataset)
+        n_rois = 264
+        fc_dim = (n_rois * (n_rois - 1)) // 2
+        X = np.zeros((n_subjects, fc_dim))
+        print(f"Created placeholder FC matrices with shape {X.shape}")
+        return X, demo_data, demo_types
+        
+    elif isinstance(dataset_or_niifiles, str):
+        # Handle real dataset with actual fMRI data
+        dataset = load_dataset(dataset_or_niifiles, split="train")
+        
+        # Load Power 264 atlas
         from nilearn import datasets
         power = datasets.fetch_coords_power_2011()
         coords = np.vstack((power.rois['x'], power.rois['y'], power.rois['z'])).T
@@ -25,52 +165,84 @@ def preprocess_fmri_to_fc(dataset_name, atlas_path=None):
             high_pass=0.01,
             t_r=2.0  # Adjust TR according to your data
         )
-    else:
-        masker = input_data.NiftiLabelsMasker(
-            labels_img=atlas_path,
+
+        # Load demographic data if needed
+        if demo_data is None:
+            if 'demographics' in dataset.features:
+                demo_df = pd.DataFrame(dataset['demographics'])
+                
+                demo_data = [
+                    demo_df['age_at_stroke'].values if 'age_at_stroke' in demo_df.columns else [],
+                    demo_df['sex'].values if 'sex' in demo_df.columns else [],
+                    demo_df['months_post_stroke'].values if 'months_post_stroke' in demo_df.columns else [],
+                    demo_df['wab_score'].values if 'wab_score' in demo_df.columns else []
+                ]
+                
+                demo_types = ['continuous', 'categorical', 'continuous', 'continuous']
+        
+        # Process fMRI data and compute FC matrices
+        fc_matrices = []
+        for nii_file in dataset['nii_files']:
+            fmri_img = load_img(nii_file)
+            time_series = masker.fit_transform(fmri_img)
+            
+            correlation_measure = connectome.ConnectivityMeasure(
+                kind='correlation', vectorize=False, discard_diagonal=False
+            )
+            
+            fc_matrix = correlation_measure.fit_transform([time_series])[0]
+            
+            triu_indices = np.triu_indices_from(fc_matrix, k=1)
+            fc_triu = fc_matrix[triu_indices]
+            
+            fc_triu = np.arctanh(fc_triu)  # Fisher z-transform
+            
+            fc_matrices.append(fc_triu)
+        
+        X = np.array(fc_matrices)
+        
+    elif isinstance(dataset_or_niifiles, list) and demo_data is not None and demo_types is not None:
+        # Handle a list of NIfTI files
+        # Similar processing as above but with local files
+        print(f"Processing {len(dataset_or_niifiles)} local NIfTI files")
+        
+        # Load Power 264 atlas
+        from nilearn import datasets
+        power = datasets.fetch_coords_power_2011()
+        coords = np.vstack((power.rois['x'], power.rois['y'], power.rois['z'])).T
+        
+        masker = input_data.NiftiSpheresMasker(
+            coords, radius=5,
             standardize=True,
             memory='nilearn_cache', memory_level=1,
             verbose=0,
             detrend=True,
             low_pass=0.1,
             high_pass=0.01,
-            t_r=2.0  # Adjust TR according to your data
-        )
-
-    # Load demographic data
-    demo_df = pd.DataFrame(dataset['demographics'])
-    
-    demo_data = [
-        demo_df['age_at_stroke'].values,
-        demo_df['sex'].values,
-        demo_df['months_post_stroke'].values,
-        demo_df['wab_score'].values
-    ]
-    
-    demo_types = ['continuous', 'categorical', 'continuous', 'continuous']
-    
-    # Process fMRI data and compute FC matrices
-    fc_matrices = []
-    for nii_file in dataset['nii_files']:
-        fmri_img = load_img(nii_file)
-        time_series = masker.fit_transform(fmri_img)
-        
-        correlation_measure = connectome.ConnectivityMeasure(
-            kind='correlation',
-            vectorize=False,
-            discard_diagonal=False
+            t_r=2.0
         )
         
-        fc_matrix = correlation_measure.fit_transform([time_series])[0]
-        
-        triu_indices = np.triu_indices_from(fc_matrix, k=1)
-        fc_triu = fc_matrix[triu_indices]
-        
-        fc_triu = np.arctanh(fc_triu)  # Fisher z-transform
-        
-        fc_matrices.append(fc_triu)
-    
-    X = np.array(fc_matrices)
+        fc_matrices = []
+        for nii_file in dataset_or_niifiles:
+            fmri_img = load_img(nii_file)
+            time_series = masker.fit_transform(fmri_img)
+            
+            correlation_measure = connectome.ConnectivityMeasure(
+                kind='correlation', vectorize=False, discard_diagonal=False
+            )
+            
+            fc_matrix = correlation_measure.fit_transform([time_series])[0]
+            
+            triu_indices = np.triu_indices_from(fc_matrix, k=1)
+            fc_triu = fc_matrix[triu_indices]
+            
+            fc_triu = np.arctanh(fc_triu)  # Fisher z-transform
+            
+            fc_matrices.append(fc_triu)
+            
+        X = np.array(fc_matrices)
+    else:
+        raise ValueError("Invalid input. Expected dataset name string or list of NIfTI files with demographic data.")
     
     # Normalize the FC data
     X = (X - np.mean(X, axis=0)) / np.std(X, axis=0)

main.py

@@ -110,63 +110,82 @@ def run_fc_analysis(data_dir="SreekarB/OSFData",
         'bsize': bsize
     })
     
-    # Load data
-    print("Loading data...")
-    nii_files, demo_data, demo_types = load_data(data_dir, demographic_file, use_hf_dataset)
-    
-    # Check if we got FC matrices directly
-    if isinstance(nii_files, list) and len(nii_files) > 0 and hasattr(nii_files[0], 'shape'):
-        print("Using pre-computed FC matrices...")
-        # Convert list of FC matrices to numpy array
-        X = np.stack([np.array(fc) for fc in nii_files])
-    else:
-        # Prepare data by converting fMRI to FC matrices
-        print("Converting fMRI data to FC matrices...")
-        X, demo_data, demo_types = preprocess_fmri_to_fc(nii_files, demo_data, demo_types)
-    
-    # Print shapes and data types
-    print(f"X shape: {X.shape}, type: {type(X)}")
-    for i, d in enumerate(demo_data):
-        print(f"Demo data {i} shape: {d.shape if hasattr(d, 'shape') else len(d)}, type: {type(d)}")
-    
-    # Train VAE and get data
-    print("Training VAE...")
-    vae, X, demo_data, demo_types = train_fc_vae(X, demo_data, demo_types, MODEL_CONFIG)
-    
-    if save_model:
-        print("Saving model...")
-        os.makedirs('models', exist_ok=True)
-        torch.save(vae.state_dict(), 'models/vae_model.pth')
-    
-    # Get latent representations
-    print("Getting latent representations...")
-    latents = vae.get_latents(X)
-    
-    # Analyze results
-    print("Analyzing demographic relationships...")
-    demographics = {
-        'age': demo_data[0],
-        'months_post_onset': demo_data[2],
-        'wab_aq': demo_data[3]
-    }
-    analysis_results = analyze_fc_patterns(latents, demographics)
-    
-    # Generate new FC matrix
-    print("Generating new FC matrices...")
-    new_demographics = [
-        [60.0],  # age
-        ['M'],   # gender
-        [12.0],  # months post onset
-        [80.0]   # wab score
-    ]
-    generated_fc = vae.transform(1, new_demographics, demo_types)
-    reconstructed_fc = vae.transform(X, demo_data, demo_types)
-    
-    # Visualize results
-    print("Creating visualizations...")
-    fig = visualize_fc_analysis(X[0], reconstructed_fc[0], generated_fc[0], analysis_results)
-    
-    return fig
+    try:
+        # Load data
+        print("Loading data...")
+        nii_files, demo_data, demo_types = load_data(data_dir, demographic_file, use_hf_dataset)
+        
+        # For SreekarB/OSFData, directly generate synthetic FC matrices
+        if data_dir == "SreekarB/OSFData" and use_hf_dataset:
+            print("Using SreekarB/OSFData dataset with synthetic FC matrices...")
+            X, demo_data, demo_types = preprocess_fmri_to_fc(data_dir, demo_data, demo_types)
+        # Check if we got FC matrices directly
+        elif isinstance(nii_files, list) and len(nii_files) > 0 and hasattr(nii_files[0], 'shape'):
+            print("Using pre-computed FC matrices...")
+            # Convert list of FC matrices to numpy array
+            X = np.stack([np.array(fc) for fc in nii_files])
+        else:
+            # Prepare data by converting fMRI to FC matrices
+            print("Converting fMRI data to FC matrices...")
+            X, demo_data, demo_types = preprocess_fmri_to_fc(nii_files, demo_data, demo_types)
+        
+        # Print shapes and data types
+        print(f"X shape: {X.shape}, type: {type(X)}")
+        for i, d in enumerate(demo_data):
+            print(f"Demo data {i} shape: {d.shape if hasattr(d, 'shape') else len(d)}, type: {type(d)}")
+        
+        # Train VAE and get data
+        print("Training VAE...")
+        vae, X, demo_data, demo_types = train_fc_vae(X, demo_data, demo_types, MODEL_CONFIG)
+        
+        if save_model:
+            print("Saving model...")
+            os.makedirs('models', exist_ok=True)
+            torch.save(vae.state_dict(), 'models/vae_model.pth')
+        
+        # Get latent representations
+        print("Getting latent representations...")
+        latents = vae.get_latents(X)
+        
+        # Analyze results
+        print("Analyzing demographic relationships...")
+        demographics = {
+            'age': demo_data[0],
+            'months_post_onset': demo_data[2],
+            'wab_aq': demo_data[3]
+        }
+        analysis_results = analyze_fc_patterns(latents, demographics)
+        
+        # Generate new FC matrix
+        print("Generating new FC matrices...")
+        new_demographics = [
+            [60.0],  # age
+            ['M'],   # gender
+            [12.0],  # months post onset
+            [80.0]   # wab score
+        ]
+        generated_fc = vae.transform(1, new_demographics, demo_types)
+        reconstructed_fc = vae.transform(X, demo_data, demo_types)
+        
+        # Visualize results
+        print("Creating visualizations...")
+        fig = visualize_fc_analysis(X[0], reconstructed_fc[0], generated_fc[0], analysis_results)
+        
+        return fig
+        
+    except Exception as e:
+        import traceback
+        print(f"Error in run_fc_analysis: {str(e)}")
+        print(traceback.format_exc())
+        
+        # Create a dummy figure with error message
+        import matplotlib.pyplot as plt
+        fig = plt.figure(figsize=(10, 6))
+        plt.text(0.5, 0.5, f"Error: {str(e)}", 
+                 horizontalalignment='center', verticalalignment='center', 
+                 fontsize=12, color='red')
+        plt.axis('off')
+        return fig
 
 if __name__ == "__main__":
     import argparse

requirements.txt

@@ -7,6 +7,6 @@ scikit-learn>=0.24.2
 matplotlib>=3.4.2
 gradio>=2.0.0
 datasets>=1.11.0
-huggingface_hub>=0.12.0
+huggingface_hub>=0.15.0
 transformers>=4.15.0