Upload 4 files

app.py

@@ -2215,7 +2215,7 @@ def create_interface():
             try:
                 # Start by setting up data for the VAE
                 from vae_model import DemoVAE
-                from data_preprocessing import load_and_preprocess_data
+                from data_preprocessing import load_and_preprocess_data, download_and_cache_dataset
                 from main import run_analysis
                 import numpy as np
                 import os
@@ -2231,36 +2231,52 @@ def create_interface():
                 logger.info("Looking for data in directory and preprocessing...")
                 
                 # This part is similar to app.train_models but only focuses on VAE
-                if data_dir == "SreekarB/OSFData1":
-                    # Use dataset, similar to existing code in app.train_models
-                    # For brevity, we'll call the full train_models function but only 
-                    # extract the VAE-related results
-                    results = app.train_models(
-                        data_dir=data_dir,
-                        latent_dim=latent_dim,
-                        nepochs=nepochs,
-                        bsize=bsize,
-                        num_samples=num_samples
-                    )
-                    
-                    # Store results in app_state for the next tabs
-                    app_state['vae'] = results.get('vae', None)
-                    app_state['latents'] = results.get('latents', None)
-                    app_state['demographics'] = results.get('demographics', None)
-                    app_state['vae_trained'] = True
-                    
-                    # Process VAE results with data preference
-                    fc_figure, learning_fig = process_vae_results(
-                        results, 
-                        prefer_real_data=prefer_real_data,
-                        use_synthetic_nifti=use_synthetic_nifti,
-                        use_synthetic_fc=use_synthetic_fc
-                    )
-                    
-                    # Return the visualizations
-                    return [fc_figure, learning_fig]
+                if data_dir == "SreekarB/OSFData1" and use_hf_dataset:
+                    # Use our improved download and cache function for HuggingFace datasets
+                    logger.info(f"Using improved HuggingFace dataset caching for {data_dir}")
+                    try:
+                        # Get the dataset and cached NIfTI files
+                        _, _, nii_files = download_and_cache_dataset(data_dir)
+                        
+                        if nii_files and len(nii_files) > 0:
+                            logger.info(f"Successfully found {len(nii_files)} NIfTI files in dataset cache")
+                            
+                            # Limit to the specified number of samples if needed
+                            if num_samples is not None and len(nii_files) > num_samples:
+                                logger.info(f"Limiting to {num_samples} NIfTI files as specified")
+                                nii_files = nii_files[:num_samples]
+                            
+                            # Use the cached files in the train_models function
+                            results = app.train_models(
+                                data_dir=data_dir,
+                                latent_dim=latent_dim,
+                                nepochs=nepochs,
+                                bsize=bsize,
+                                num_samples=num_samples,
+                                hf_nii_files=nii_files  # Use our cached files
+                            )
+                        else:
+                            # Fallback to standard train_models
+                            logger.warning("No NIfTI files found in cache, using standard train_models")
+                            results = app.train_models(
+                                data_dir=data_dir,
+                                latent_dim=latent_dim,
+                                nepochs=nepochs,
+                                bsize=bsize,
+                                num_samples=num_samples
+                            )
+                    except Exception as cache_err:
+                        logger.error(f"Error using cached dataset: {cache_err}")
+                        # Fallback to standard train_models
+                        results = app.train_models(
+                            data_dir=data_dir,
+                            latent_dim=latent_dim,
+                            nepochs=nepochs,
+                            bsize=bsize,
+                            num_samples=num_samples
+                        )
                 else:
-                    # Local directory case
+                    # Local directory case or non-HuggingFace dataset
                     results = app.train_models(
                         data_dir=data_dir,
                         latent_dim=latent_dim,
@@ -2268,23 +2284,23 @@ def create_interface():
                         bsize=bsize,
                         num_samples=num_samples
                     )
-                    
-                    # Store results in app_state
-                    app_state['vae'] = results.get('vae', None)
-                    app_state['latents'] = results.get('latents', None)
-                    app_state['demographics'] = results.get('demographics', None)
-                    app_state['vae_trained'] = True
-                    
-                    # Process VAE results with data preference
-                    fc_figure, learning_fig = process_vae_results(
-                        results, 
-                        prefer_real_data=prefer_real_data,
-                        use_synthetic_nifti=use_synthetic_nifti,
-                        use_synthetic_fc=use_synthetic_fc
-                    )
-                    
-                    # Return the visualizations
-                    return [fc_figure, learning_fig]
+                
+                # Store results in app_state
+                app_state['vae'] = results.get('vae', None)
+                app_state['latents'] = results.get('latents', None)
+                app_state['demographics'] = results.get('demographics', None)
+                app_state['vae_trained'] = True
+                
+                # Process VAE results with data preference
+                fc_figure, learning_fig = process_vae_results(
+                    results, 
+                    prefer_real_data=prefer_real_data,
+                    use_synthetic_nifti=use_synthetic_nifti,
+                    use_synthetic_fc=use_synthetic_fc
+                )
+                
+                # Return the visualizations
+                return [fc_figure, learning_fig]
             except Exception as e:
                 logger.error(f"Error in VAE training: {str(e)}", exc_info=True)
                 error_fig = plt.figure(figsize=(10, 6))

data_preprocessing.py

@@ -137,6 +137,9 @@ def process_single_fmri(fmri_file, allow_synthetic=False, use_cache=True, try_pr
     """
     print(f"Processing fMRI file: {fmri_file}")
     
+    # Make sure os is imported to avoid reference error
+    import os
+    
     # Check if cached FC matrix exists
     if use_cache:
         file_hash = get_file_hash(fmri_file)
@@ -334,6 +337,7 @@ def process_single_fmri(fmri_file, allow_synthetic=False, use_cache=True, try_pr
             
             # Import nilearn preprocessing
             from nilearn import image
+            import os  # Ensure os is imported again here
             
             # Load the image
             orig_img = load_img(fmri_file)
@@ -346,7 +350,6 @@ def process_single_fmri(fmri_file, allow_synthetic=False, use_cache=True, try_pr
             
             # Create a temporary preprocessed file
             import tempfile
-            import os
             
             # Create a preprocessing directory
             preproc_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'preproc')
@@ -742,6 +745,47 @@ def clear_cache(cache_type=None):
     
     print("Cache clearing complete")
 
+def download_and_cache_dataset(dataset_name, cache_dir=None):
+    """
+    Download a dataset from HuggingFace and save it to a local cache.
+    
+    Args:
+        dataset_name (str): Name of the dataset on HuggingFace (e.g., 'SreekarB/OSFData1')
+        cache_dir (str, optional): Directory to store the cached files. If None, uses the default HuggingFace cache.
+    
+    Returns:
+        dataset: The loaded dataset object
+        cache_path (str): Path to the cached dataset
+        nii_files (list): List of NIfTI file paths found in the dataset
+    """
+    from datasets import load_dataset
+    import os
+    from huggingface_hub import HfFolder
+    
+    print(f"Loading dataset: {dataset_name}")
+    
+    # Load the dataset (this will download it if not already cached)
+    dataset = load_dataset(dataset_name, cache_dir=cache_dir)
+    
+    # Get the cache directory
+    if cache_dir is None:
+        cache_dir = os.path.join(HfFolder.get_cache_dir(), "datasets")
+    
+    # The dataset is now cached locally
+    dataset_cache_path = os.path.join(cache_dir, dataset_name.replace("/", "--"))
+    print(f"Dataset cached at: {dataset_cache_path}")
+    
+    # Locate NIfTI files in the cached dataset
+    nii_files = []
+    for root, dirs, filenames in os.walk(dataset_cache_path):
+        for filename in filenames:
+            if filename.endswith('.nii') or filename.endswith('.nii.gz'):
+                nii_files.append(os.path.join(root, filename))
+    
+    print(f"Found {len(nii_files)} NIfTI files in dataset cache")
+    
+    return dataset, dataset_cache_path, nii_files
+
 def load_and_preprocess_data(data_dir, demographic_file, use_hf_dataset=False, 
                         hf_nii_files=None, hf_demo_data=None, hf_demo_types=None,
                         max_samples=None):
@@ -796,6 +840,20 @@ def load_and_preprocess_data(data_dir, demographic_file, use_hf_dataset=False,
             else:
                 # No demographic data provided
                 raise ValueError("No demographic data provided")
+        
+        # Try to properly access the HuggingFace dataset files
+        if not (hf_nii_files and len(hf_nii_files) > 0):
+            try:
+                print(f"Attempting to download and cache dataset: {data_dir}")
+                # Try to load and cache the dataset files
+                _, _, nii_files_from_hf = download_and_cache_dataset(data_dir)
+                
+                if nii_files_from_hf and len(nii_files_from_hf) > 0:
+                    print(f"Successfully found {len(nii_files_from_hf)} NIfTI files in HuggingFace dataset cache")
+                    hf_nii_files = nii_files_from_hf
+            except Exception as e:
+                print(f"Error accessing HuggingFace dataset files: {e}")
+                hf_nii_files = []
                 
         # Use provided NIfTI files from HuggingFace
         if hf_nii_files and len(hf_nii_files) > 0: