Upload 9 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
 ---

analysis.py

@@ -1,16 +1,39 @@
 from scipy.stats import pearsonr
+import numpy as np
+import warnings
 
 def analyze_fc_patterns(latents, demographics):
     results = {}
-    for demo_name, demo_values in demographics.items():
-        if demo_name != 'sex':  # For continuous variables
-            correlations = []
-            p_values = []
-            for latent_dim in range(latents.shape[1]):
-                r, p = pearsonr(latents[:, latent_dim], demo_values)
-                correlations.append(r)
-                p_values.append(p)
-            results[demo_name] = {'correlations': correlations, 'p_values': p_values}
+    
+    # Suppress the ConstantInputWarning
+    with warnings.catch_warnings():
+        warnings.filterwarnings('ignore', category=RuntimeWarning)
+        
+        for demo_name, demo_values in demographics.items():
+            # Check if the demographic is categorical or continuous
+            if demo_name not in ['sex', 'gender']:  # For continuous variables
+                correlations = []
+                p_values = []
+                
+                for latent_dim in range(latents.shape[1]):
+                    # Check for constant values that would cause correlation issues
+                    if np.all(latents[:, latent_dim] == latents[0, latent_dim]) or np.all(demo_values == demo_values[0]):
+                        # If either array is constant, correlation is undefined
+                        r, p = np.nan, np.nan
+                    else:
+                        try:
+                            # Convert to numpy array if not already
+                            demo_array = np.array(demo_values, dtype=float)
+                            # Calculate correlation
+                            r, p = pearsonr(latents[:, latent_dim], demo_array)
+                        except (ValueError, TypeError) as e:
+                            print(f"Error calculating correlation for {demo_name}, dimension {latent_dim}: {e}")
+                            r, p = np.nan, np.nan
+                    
+                    correlations.append(r)
+                    p_values.append(p)
+                
+                results[demo_name] = {'correlations': correlations, 'p_values': p_values}
     
     return results
 

data_preprocessing.py

@@ -23,9 +23,40 @@ def preprocess_fmri_to_fc(dataset_or_niifiles, demo_data=None, demo_types=None):
     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")
+    if isinstance(dataset_or_niifiles, str):
+        dataset_name = dataset_or_niifiles
+        print(f"Loading data from dataset: {dataset_name}")
+        try:
+            # Try multiple approaches to load the dataset
+            approaches = [
+                lambda: load_dataset(dataset_name, split="train"),
+                lambda: load_dataset(dataset_name),  # Try without split
+                lambda: load_dataset(dataset_name, split="train", trust_remote_code=True),  # Try with trust_remote_code
+                lambda: load_dataset(dataset_name.split("/")[-1], split="train") if "/" in dataset_name else None
+            ]
+            
+            dataset = None
+            last_error = None
+            
+            for i, approach in enumerate(approaches):
+                if approach is None:
+                    continue
+                    
+                try:
+                    print(f"Attempt {i+1} to load dataset...")
+                    dataset = approach()
+                    print(f"Successfully loaded dataset with approach {i+1}!")
+                    break
+                except Exception as e:
+                    print(f"Attempt {i+1} failed: {e}")
+                    last_error = e
+                    
+            if dataset is None:
+                print(f"All attempts to load dataset failed. Last error: {last_error}")
+                raise ValueError(f"Could not load dataset {dataset_name}")
+        except Exception as e:
+            print(f"Error during dataset loading: {e}")
+            raise
         
         # Prepare demographics data from the dataset
         if demo_data is None:
@@ -50,93 +81,369 @@ def preprocess_fmri_to_fc(dataset_or_niifiles, demo_data=None, demo_types=None):
         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")
+        # Create a temp directory for downloads
+        import tempfile
+        from huggingface_hub import hf_hub_download
+        import shutil
+        
+        temp_dir = tempfile.mkdtemp(prefix="hf_nifti_")
+        print(f"Created temporary directory for NIfTI files: {temp_dir}")
+        
+        try:
+            # First approach: Check if there are any columns containing file paths
+            nii_columns = []
+            for col in dataset.column_names:
+                # Check if column name suggests NIfTI files
+                if 'nii' in col.lower() or 'nifti' in col.lower() or 'fmri' in col.lower():
+                    nii_columns.append(col)
+                # Or check if column contains file paths
+                elif len(dataset) > 0:
+                    first_val = dataset[0][col]
+                    if isinstance(first_val, str) and (first_val.endswith('.nii') or first_val.endswith('.nii.gz')):
+                        nii_columns.append(col)
+            
+            if nii_columns:
+                print(f"Found columns that may contain NIfTI files: {nii_columns}")
                 
-                # 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
-                    )
+                for col in nii_columns:
+                    print(f"Processing column '{col}'...")
                     
-                    # 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)
+                    for i, item in enumerate(dataset[col]):
+                        if not isinstance(item, str):
+                            print(f"Item {i} in column {col} is not a string but {type(item)}")
+                            continue
+                            
+                        if not (item.endswith('.nii') or item.endswith('.nii.gz')):
+                            print(f"Item {i} in column {col} is not a NIfTI file: {item}")
+                            continue
                             
-                            correlation_measure = connectome.ConnectivityMeasure(
-                                kind='correlation',
-                                vectorize=False,
-                                discard_diagonal=False
+                        print(f"Downloading {item} from dataset {dataset_name}...")
+                        
+                        try:
+                            # Attempt to download with explicit filename
+                            file_path = hf_hub_download(
+                                repo_id=dataset_name,
+                                filename=item,
+                                repo_type="dataset",
+                                cache_dir=temp_dir
                             )
+                            nii_files.append(file_path)
+                            print(f"✓ Successfully downloaded {item}")
+                        except Exception as e1:
+                            print(f"Error downloading with explicit filename: {e1}")
                             
-                            fc_matrix = correlation_measure.fit_transform([time_series])[0]
+                            # Second attempt: try with the item's basename
+                            try:
+                                basename = os.path.basename(item)
+                                print(f"Trying with basename: {basename}")
+                                file_path = hf_hub_download(
+                                    repo_id=dataset_name,
+                                    filename=basename,
+                                    repo_type="dataset",
+                                    cache_dir=temp_dir
+                                )
+                                nii_files.append(file_path)
+                                print(f"✓ Successfully downloaded {basename}")
+                            except Exception as e2:
+                                print(f"Error downloading with basename: {e2}")
+                                
+                                # Third attempt: check if it's a binary blob in the dataset
+                                try:
+                                    if hasattr(dataset[i], 'keys') and 'bytes' in dataset[i]:
+                                        print("Found binary data in dataset, saving to temporary file...")
+                                        binary_data = dataset[i]['bytes']
+                                        temp_file = os.path.join(temp_dir, basename)
+                                        with open(temp_file, 'wb') as f:
+                                            f.write(binary_data)
+                                        nii_files.append(temp_file)
+                                        print(f"✓ Saved binary data to {temp_file}")
+                                except Exception as e3:
+                                    print(f"Error handling binary data: {e3}")
+                                    
+                                # Last resort: look 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} anywhere")
+            
+            # Second approach: Try to find NIfTI files in dataset repository directly
+            if not nii_files:
+                print("No NIfTI files found in dataset columns. Trying direct repository search...")
+                
+                try:
+                    from huggingface_hub import list_repo_files, hf_hub_download
+                    
+                    # Try to list all files in the repository
+                    try:
+                        print("Listing all repository files...")
+                        all_repo_files = list_repo_files(dataset_name, repo_type="dataset")
+                        print(f"Found {len(all_repo_files)} files in repository")
+                        
+                        # First prioritize P*_rs.nii files
+                        p_rs_files = [f for f in all_repo_files if f.endswith('_rs.nii') and f.startswith('P')]
+                        
+                        # Then include all other NIfTI files
+                        other_nii_files = [f for f in all_repo_files if (f.endswith('.nii') or f.endswith('.nii.gz')) and f not in p_rs_files]
+                        
+                        # Combine, with P*_rs.nii files first
+                        nii_repo_files = p_rs_files + other_nii_files
+                        
+                        if nii_repo_files:
+                            print(f"Found {len(nii_repo_files)} NIfTI files in repository: {nii_repo_files[:5] if len(nii_repo_files) > 5 else nii_repo_files}...")
                             
-                            triu_indices = np.triu_indices_from(fc_matrix, k=1)
-                            fc_triu = fc_matrix[triu_indices]
+                            # Download each file
+                            for nii_file in nii_repo_files:
+                                try:
+                                    file_path = hf_hub_download(
+                                        repo_id=dataset_name,
+                                        filename=nii_file,
+                                        repo_type="dataset",
+                                        cache_dir=temp_dir
+                                    )
+                                    nii_files.append(file_path)
+                                    print(f"✓ Downloaded {nii_file}")
+                                except Exception as e:
+                                    print(f"Error downloading {nii_file}: {e}")
+                    except Exception as e:
+                        print(f"Error listing repository files: {e}")
+                        print("Will try alternative approaches...")
+                    
+                    # If repo listing fails, try with common NIfTI file patterns directly
+                    if not nii_files:
+                        print("Trying common NIfTI file patterns...")
+                        
+                        # Focus specifically on P*_rs.nii pattern
+                        patterns = []
+                        
+                        # Generate P01_rs.nii through P30_rs.nii
+                        for i in range(1, 31):  # Try subjects 1-30
+                            patterns.append(f"P{i:02d}_rs.nii")
                             
-                            fc_triu = np.arctanh(fc_triu)  # Fisher z-transform
+                        # Also try with .nii.gz extension
+                        for i in range(1, 31):
+                            patterns.append(f"P{i:02d}_rs.nii.gz")
                             
-                            fc_matrices.append(fc_triu)
-                            print(f"Processed {nii_file} to FC matrix")
+                        # Include a few other common patterns as fallbacks
+                        patterns.extend([
+                            "sub-01_task-rest_bold.nii.gz",  # BIDS format
+                            "fmri.nii.gz", "bold.nii.gz", 
+                            "rest.nii.gz"
+                        ])
+                        
+                        for pattern in patterns:
+                            try:
+                                print(f"Trying to download {pattern}...")
+                                file_path = hf_hub_download(
+                                    repo_id=dataset_name,
+                                    filename=pattern,
+                                    repo_type="dataset",
+                                    cache_dir=temp_dir
+                                )
+                                nii_files.append(file_path)
+                                print(f"✓ Successfully downloaded {pattern}")
+                            except Exception as e:
+                                print(f"× Failed to download {pattern}")
+                                
+                    # If we still couldn't find any files, check if data files are nested
+                    if not nii_files:
+                        print("Checking for nested data files...")
+                        nested_paths = ["data/", "raw/", "nii/", "derivatives/", "fmri/", "nifti/"]
+                        
+                        for path in nested_paths:
+                            for pattern in patterns:
+                                nested_file = f"{path}{pattern}"
+                                try:
+                                    print(f"Trying to download {nested_file}...")
+                                    file_path = hf_hub_download(
+                                        repo_id=dataset_name,
+                                        filename=nested_file,
+                                        repo_type="dataset",
+                                        cache_dir=temp_dir
+                                    )
+                                    nii_files.append(file_path)
+                                    print(f"✓ Successfully downloaded {nested_file}")
+                                    # If we found one file in this directory, try to find all files in it
+                                    try:
+                                        all_files_in_dir = [f for f in all_repo_files if f.startswith(path)]
+                                        nii_files_in_dir = [f for f in all_files_in_dir if f.endswith('.nii') or f.endswith('.nii.gz')]
+                                        print(f"Found {len(nii_files_in_dir)} additional NIfTI files in {path}")
+                                        
+                                        for nii_file in nii_files_in_dir:
+                                            if nii_file != nested_file:  # Skip the one we already downloaded
+                                                try:
+                                                    file_path = hf_hub_download(
+                                                        repo_id=dataset_name,
+                                                        filename=nii_file,
+                                                        repo_type="dataset",
+                                                        cache_dir=temp_dir
+                                                    )
+                                                    nii_files.append(file_path)
+                                                    print(f"✓ Downloaded {nii_file}")
+                                                except Exception as e:
+                                                    print(f"Error downloading {nii_file}: {e}")
+                                    except Exception as e:
+                                        print(f"Error finding additional files in {path}: {e}")
+                                except Exception as e:
+                                    pass
+                                
+                except Exception as e:
+                    print(f"Error during repository exploration: {e}")
+            
+            # If we still don't have any files, try to search for P*_rs.nii pattern specifically
+            if not nii_files:
+                print("Trying to find files matching P*_rs.nii pattern specifically...")
+                
+                try:
+                    # List all files in the repository (if we haven't already)
+                    if not 'all_repo_files' in locals():
+                        from huggingface_hub import list_repo_files
+                        try:
+                            all_repo_files = list_repo_files(dataset_name, repo_type="dataset")
                         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(f"Error listing repo files: {e}")
+                            all_repo_files = []
+                    
+                    # Look for files matching the pattern exactly (P*_rs.nii)
+                    pattern_files = [f for f in all_repo_files if '_rs.nii' in f and f.startswith('P')]
+                    
+                    # If we don't find any exact matches, try a more relaxed pattern
+                    if not pattern_files:
+                        pattern_files = [f for f in all_repo_files if 'rs.nii' in f.lower()]
+                    
+                    if pattern_files:
+                        print(f"Found {len(pattern_files)} files matching rs.nii pattern")
+                        
+                        # Download each file
+                        for pattern_file in pattern_files:
+                            try:
+                                file_path = hf_hub_download(
+                                    repo_id=dataset_name,
+                                    filename=pattern_file,
+                                    repo_type="dataset",
+                                    cache_dir=temp_dir
+                                )
+                                nii_files.append(file_path)
+                                print(f"✓ Downloaded {pattern_file}")
+                            except Exception as e:
+                                print(f"Error downloading {pattern_file}: {e}")
+                except Exception as e:
+                    print(f"Error searching for pattern files: {e}")
+                
+            print(f"Found total of {len(nii_files)} NIfTI files")
+        except Exception as e:
+            print(f"Unexpected error during NIfTI file search: {e}")
+            import traceback
+            traceback.print_exc()
         
+        # 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 = []
+            valid_files = 0
+            total_files = len(nii_files)
+            
+            for nii_file in nii_files:
+                try:
+                    print(f"Processing {nii_file}...")
+                    fmri_img = load_img(nii_file)
+                    
+                    # Check image dimensions
+                    if len(fmri_img.shape) < 4 or fmri_img.shape[3] < 10:
+                        print(f"Warning: {nii_file} has insufficient time points: {fmri_img.shape}")
+                        continue
+                        
+                    time_series = masker.fit_transform(fmri_img)
+                    
+                    # Validate time series data
+                    if np.isnan(time_series).any() or np.isinf(time_series).any():
+                        print(f"Warning: {nii_file} contains NaN or Inf values after masking")
+                        # Replace NaNs with zeros for this file
+                        time_series = np.nan_to_num(time_series)
+                    
+                    correlation_measure = connectome.ConnectivityMeasure(
+                        kind='correlation',
+                        vectorize=False,
+                        discard_diagonal=False
+                    )
+                    
+                    fc_matrix = correlation_measure.fit_transform([time_series])[0]
+                    
+                    # Check for invalid correlation values
+                    if np.isnan(fc_matrix).any():
+                        print(f"Warning: {nii_file} produced NaN correlation values")
+                        continue
+                        
+                    triu_indices = np.triu_indices_from(fc_matrix, k=1)
+                    fc_triu = fc_matrix[triu_indices]
+                    
+                    # Fisher z-transform with proper bounds check
+                    # Clip correlation values to valid range for arctanh
+                    fc_triu_clipped = np.clip(fc_triu, -0.999, 0.999)
+                    fc_triu = np.arctanh(fc_triu_clipped)
+                    
+                    fc_matrices.append(fc_triu)
+                    valid_files += 1
+                    print(f"Successfully processed {nii_file} to FC matrix")
+                    
+                except Exception as e:
+                    print(f"Error processing {nii_file}: {e}")
+            
+            if fc_matrices:
+                print(f"Successfully processed {valid_files} out of {total_files} files")
+                
+                # Ensure all matrices have the same dimensions
+                dims = [m.shape[0] for m in fc_matrices]
+                if len(set(dims)) > 1:
+                    print(f"Warning: FC matrices have inconsistent dimensions: {dims}")
+                    # Use the most common dimension
+                    from collections import Counter
+                    most_common_dim = Counter(dims).most_common(1)[0][0]
+                    print(f"Using most common dimension: {most_common_dim}")
+                    fc_matrices = [m for m in fc_matrices if m.shape[0] == most_common_dim]
+                
+                X = np.array(fc_matrices)
+                
+                # Normalize the FC data
+                mean_x = np.mean(X, axis=0)
+                std_x = np.std(X, axis=0)
+                
+                # Handle zero standard deviation
+                std_x[std_x == 0] = 1.0
+                
+                X = (X - mean_x) / std_x
+                print(f"Created FC matrices with shape {X.shape}")
+                
+                # Make sure demo_data matches the number of FC matrices
+                if len(demo_data[0]) != X.shape[0]:
+                    print(f"Warning: Number of subjects in demographic data ({len(demo_data[0])}) " +
+                          f"doesn't match number of FC matrices ({X.shape[0]})")
+                    # Adjust demo_data to match FC matrices
+                    indices = list(range(min(len(demo_data[0]), X.shape[0])))
+                    X = X[indices]
+                    demo_data = [d[indices] for d in demo_data]
+                    
+                return X, demo_data, demo_types
+
         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)
@@ -247,5 +554,4 @@ def preprocess_fmri_to_fc(dataset_or_niifiles, demo_data=None, demo_types=None):
     # Normalize the FC data
     X = (X - np.mean(X, axis=0)) / np.std(X, axis=0)
     
-    return X, demo_data, demo_types
-
+    return X, demo_data, demo_types

main.py

@@ -25,6 +25,23 @@ def train_fc_vae(X, demo_data, demo_types, model_config):
     
     print(f"Creating VAE with latent dim={model_config['latent_dim']}, epochs={model_config['nepochs']}")
     
+    # Ensure X is a numpy array with correct data type
+    if not isinstance(X, np.ndarray):
+        print(f"Converting X from {type(X)} to numpy array")
+        X = np.array(X, dtype=np.float32)
+    
+    # Ensure demo_data contains numpy arrays
+    for i, d in enumerate(demo_data):
+        if not isinstance(d, np.ndarray):
+            print(f"Converting demographic {i} from {type(d)} to numpy array")
+            demo_data[i] = np.array(d)
+    
+    # Check for NaN or Inf values
+    if np.isnan(X).any() or np.isinf(X).any():
+        print("Warning: X contains NaN or Inf values. Replacing with zeros.")
+        X = np.nan_to_num(X)
+    
+    # Create the VAE model
     vae = DemoVAE(
         latent_dim=model_config['latent_dim'],
         nepochs=model_config['nepochs'],
@@ -194,16 +211,32 @@ def run_fc_analysis(data_dir="SreekarB/OSFData",
         # Generate new FC matrix
         print("Generating new FC matrices...")
         
+        # Get data types from original demographic data for proper conversion
+        demo_dtypes = [type(d[0]) if len(d) > 0 else float for d in demo_data]
+        
         # Convert to numpy arrays to avoid "expected np.ndarray (got list)" error
         new_demographics = [
-            np.array([60.0]),  # age
-            np.array(['M']),   # gender
-            np.array([12.0]),  # months post onset
-            np.array([80.0])   # wab score
+            np.array([60.0], dtype=np.float64),        # age
+            np.array(['M'], dtype=np.str_),           # gender
+            np.array([12.0], dtype=np.float64),       # months post onset
+            np.array([80.0], dtype=np.float64)        # wab score
         ]
         
+        # Verify the demographic data arrays match the expected types
+        print("Demographic data types:")
+        for i, (name, data) in enumerate(zip(['age', 'gender', 'mpo', 'wab'], new_demographics)):
+            print(f"  {name}: shape={data.shape}, dtype={data.dtype}")
+        
         print("Generating FC matrix with demographic values: age=60, gender=M, mpo=12, wab=80")
-        generated_fc = vae.transform(1, new_demographics, demo_types)
+        try:
+            generated_fc = vae.transform(1, new_demographics, demo_types)
+        except Exception as e:
+            print(f"Error generating new FC matrix: {e}")
+            # Try with a fallback approach
+            print("Trying alternative generation approach...")
+            # If specific gender is causing issues, try the first gender from training data
+            new_demographics[1] = np.array([demo_data[1][0]])
+            generated_fc = vae.transform(1, new_demographics, demo_types)
         reconstructed_fc = vae.transform(X, demo_data, demo_types)
         
         # Visualize results

src/demovae/model.py

@@ -14,7 +14,11 @@ def to_torch(x):
 
 def to_cuda(x, use_cuda):
     if use_cuda:
-        return x.cuda()
+        try:
+            return x.cuda()
+        except (RuntimeError, AssertionError) as e:
+            print(f"Warning: CUDA error: {e}. Falling back to CPU.")
+            return x
     else:
         return x
 

src/demovae/sklearn.py

@@ -1,5 +1,6 @@
 
-from demovae.model import VAE, train_vae, to_torch, to_cuda, to_numpy, demo_to_torch
+from .model import VAE, train_vae, to_torch, to_cuda, to_numpy, demo_to_torch
+import numpy as np
 
 from sklearn.base import BaseEstimator
 

test_hf_download.py

@@ -0,0 +1,237 @@
+import os
+import sys
+import tempfile
+from huggingface_hub import hf_hub_download, list_repo_files, HfApi
+from datasets import load_dataset
+import numpy as np
+import pandas as pd
+
+def test_huggingface_download(dataset_name="SreekarB/OSFData", revision=None, auth_token=None):
+    """
+    Test script to verify downloading NIfTI files from HuggingFace Datasets
+    """
+    print(f"Testing download from HuggingFace dataset: {dataset_name}")
+    
+    # Create a temporary directory for downloads
+    temp_dir = tempfile.mkdtemp(prefix="hf_test_")
+    print(f"Created temporary directory: {temp_dir}")
+    
+    try:
+        # Step 1: Load the dataset and check its structure
+        print("\n1. Loading dataset...")
+        try:
+            print(f"Attempting to load dataset: {dataset_name}, revision: {revision}, with auth token: {'Yes' if auth_token else 'No'}")
+            
+            # Try multiple approaches to load the dataset
+            approaches = [
+                lambda: load_dataset(dataset_name, split="train", revision=revision, token=auth_token),
+                lambda: load_dataset(dataset_name, revision=revision, token=auth_token),  # Try without split
+                lambda: load_dataset(dataset_name, split="train", trust_remote_code=True, revision=revision, token=auth_token),  # Try with trust_remote_code
+                # If the dataset name has a slash, try just the second part
+                lambda: load_dataset(dataset_name.split("/")[-1], split="train", revision=revision, token=auth_token) if "/" in dataset_name else None
+            ]
+            
+            dataset = None
+            last_error = None
+            
+            for i, approach in enumerate(approaches):
+                if approach is None:
+                    continue
+                
+                try:
+                    print(f"Attempt {i+1}...")
+                    dataset = approach()
+                    print(f"Attempt {i+1} succeeded!")
+                    break
+                except Exception as e:
+                    print(f"Attempt {i+1} failed: {e}")
+                    last_error = e
+            
+            if dataset is None:
+                print(f"All attempts to load dataset failed. Last error: {last_error}")
+                
+                # Try direct API approach
+                print("Trying direct HF API approach...")
+                try:
+                    api = HfApi()
+                    repo_info = api.repo_info(repo_id=dataset_name, repo_type="dataset")
+                    print(f"Dataset exists on HF: {repo_info.id}")
+                    print(f"Dataset info: Private: {repo_info.private}, Size: {repo_info.size}")
+                except Exception as e:
+                    print(f"API check failed: {e}")
+                return
+            
+            print(f"Successfully loaded dataset with {len(dataset)} items")
+            print(f"Dataset structure: {type(dataset)}")
+            
+            try:
+                print(f"Dataset columns: {dataset.column_names}")
+                if len(dataset) > 0:
+                    print(f"First item keys: {list(dataset[0].keys())}")
+            except Exception as e:
+                print(f"Error accessing dataset structure: {e}")
+        except Exception as e:
+            print(f"Error loading dataset: {e}")
+            return
+            
+        # Step 2: Look for columns that might contain NIfTI files
+        print("\n2. Searching for columns with NIfTI files...")
+        nii_columns = []
+        for col in dataset.column_names:
+            if 'nii' in col.lower() or 'nifti' in col.lower() or 'fmri' in col.lower():
+                nii_columns.append(col)
+                continue
+                
+            if len(dataset) > 0:
+                try:
+                    first_val = dataset[0][col]
+                    if isinstance(first_val, str) and (first_val.endswith('.nii') or first_val.endswith('.nii.gz')):
+                        nii_columns.append(col)
+                except:
+                    pass
+                    
+        if nii_columns:
+            print(f"Found {len(nii_columns)} columns that may contain NIfTI files: {nii_columns}")
+            
+            # Step 3: Try to download a file from each column
+            print("\n3. Attempting to download from columns...")
+            downloaded = False
+            
+            for col in nii_columns:
+                print(f"\nTrying column '{col}'...")
+                
+                # Get the first non-empty value
+                file_path = None
+                for i, item in enumerate(dataset[col]):
+                    if isinstance(item, str) and (item.endswith('.nii') or item.endswith('.nii.gz')):
+                        file_path = item
+                        print(f"Found NIfTI file path at index {i}: {file_path}")
+                        break
+                
+                if not file_path:
+                    print(f"No valid NIfTI file paths found in column {col}")
+                    continue
+                    
+                # Try to download the file
+                try:
+                    downloaded_path = hf_hub_download(
+                        repo_id=dataset_name,
+                        filename=file_path,
+                        repo_type="dataset",
+                        cache_dir=temp_dir
+                    )
+                    print(f"✓ Successfully downloaded to: {downloaded_path}")
+                    downloaded = True
+                    break
+                except Exception as e:
+                    print(f"× Failed to download {file_path}: {e}")
+                    
+                    # Try with basename
+                    try:
+                        basename = os.path.basename(file_path)
+                        print(f"Trying with basename: {basename}")
+                        downloaded_path = hf_hub_download(
+                            repo_id=dataset_name,
+                            filename=basename,
+                            repo_type="dataset",
+                            cache_dir=temp_dir
+                        )
+                        print(f"✓ Successfully downloaded to: {downloaded_path}")
+                        downloaded = True
+                        break
+                    except Exception as e:
+                        print(f"× Failed to download {basename}: {e}")
+            
+            if downloaded:
+                print("\nSuccessfully downloaded a file from column data!")
+            else:
+                print("\nFailed to download any files from columns.")
+        else:
+            print("No columns found that might contain NIfTI files.")
+            
+        # Step 4: Try to list and download from repository directly
+        print("\n4. Searching repository files directly...")
+        try:
+            all_files = list_repo_files(dataset_name, repo_type="dataset")
+            print(f"Found {len(all_files)} files in repository")
+            
+            # Look for NIfTI files
+            nii_files = [f for f in all_files if f.endswith('.nii') or f.endswith('.nii.gz')]
+            rs_files = [f for f in all_files if 'rs.nii' in f.lower()]
+            
+            print(f"Found {len(nii_files)} .nii/.nii.gz files")
+            print(f"Found {len(rs_files)} files matching 'rs.nii' pattern")
+            
+            if nii_files:
+                # Try to download the first file
+                test_file = nii_files[0]
+                print(f"Attempting to download: {test_file}")
+                
+                try:
+                    downloaded_path = hf_hub_download(
+                        repo_id=dataset_name,
+                        filename=test_file,
+                        repo_type="dataset",
+                        cache_dir=temp_dir
+                    )
+                    print(f"✓ Successfully downloaded to: {downloaded_path}")
+                except Exception as e:
+                    print(f"× Failed to download {test_file}: {e}")
+            else:
+                print("No NIfTI files found in repository listing.")
+                
+            # Step 5: Try any P01_rs.nii pattern files specifically
+            if rs_files:
+                print("\n5. Trying P01_rs.nii pattern files...")
+                test_file = rs_files[0]
+                print(f"Attempting to download: {test_file}")
+                
+                try:
+                    downloaded_path = hf_hub_download(
+                        repo_id=dataset_name,
+                        filename=test_file,
+                        repo_type="dataset",
+                        cache_dir=temp_dir
+                    )
+                    print(f"✓ Successfully downloaded to: {downloaded_path}")
+                except Exception as e:
+                    print(f"× Failed to download {test_file}: {e}")
+            
+        except Exception as e:
+            print(f"Error listing repository files: {e}")
+            
+        # Step 6: Check if dataset is accessible through HF API
+        print("\n6. Checking dataset through HF API...")
+        try:
+            api = HfApi()
+            repo_info = api.repo_info(repo_id=dataset_name, repo_type="dataset")
+            print(f"Repository info: {repo_info.sha}, {repo_info.lastModified}")
+            
+            # Check if there are binary files
+            lfs_files = [f for f in all_files if 'lfs' in api.get_paths_info(dataset_name, paths=[f]).paths[0].lfs]
+            print(f"Found {len(lfs_files)} LFS (potentially binary) files")
+            if lfs_files and len(lfs_files) > 0:
+                print(f"First LFS file: {lfs_files[0]}")
+        except Exception as e:
+            print(f"Error accessing dataset through API: {e}")
+            
+    except Exception as e:
+        print(f"Unexpected error during testing: {e}")
+        import traceback
+        traceback.print_exc()
+    finally:
+        print(f"\nTest completed. Temporary directory: {temp_dir}")
+        # Uncomment to clean up: shutil.rmtree(temp_dir)
+
+if __name__ == "__main__":
+    # Process command line arguments
+    import argparse
+    parser = argparse.ArgumentParser(description='Test HuggingFace dataset downloading')
+    parser.add_argument('--dataset', type=str, default="SreekarB/OSFData", help='HuggingFace dataset name')
+    parser.add_argument('--revision', type=str, default=None, help='Dataset revision/branch')
+    parser.add_argument('--token', type=str, default=None, help='HuggingFace authentication token')
+    
+    args = parser.parse_args()
+    
+    # Use command line arguments
+    test_huggingface_download(args.dataset, args.revision, args.token)