|
|
import os |
|
|
import glob |
|
|
import numpy as np |
|
|
from typing import Any, Callable, Dict, Optional, Set, Tuple |
|
|
import torch |
|
|
from torch.utils.data import Dataset |
|
|
import random |
|
|
|
|
|
class fMRIDataset(Dataset): |
|
|
def __init__(self, |
|
|
data_root, datasets, split_suffixes, crop_length=40, downstream=False): |
|
|
|
|
|
self.file_paths = [] |
|
|
self.crop_length = crop_length |
|
|
self.downstream = downstream |
|
|
for dataset_name in datasets: |
|
|
for suffix in split_suffixes: |
|
|
folder_name = f"{dataset_name}_{suffix}" |
|
|
folder_path = os.path.join(data_root, folder_name) |
|
|
if not os.path.exists(folder_path): |
|
|
print(f"Warning: Folder not found: {folder_path}") |
|
|
continue |
|
|
|
|
|
for root, dirs, files in os.walk(folder_path): |
|
|
npz_files = glob.glob(os.path.join(root, "*.npz")) |
|
|
if len(npz_files) > 1: |
|
|
|
|
|
|
|
|
npz_files = sorted(npz_files)[:1] |
|
|
self.file_paths.extend(npz_files) |
|
|
|
|
|
print(f"Dataset loaded. Total files found: {len(self.file_paths)}") |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.file_paths) |
|
|
|
|
|
def __getitem__(self, idx): |
|
|
|
|
|
file_path = self.file_paths[idx] |
|
|
try: |
|
|
with np.load(file_path) as data_file: |
|
|
key = list(data_file.keys())[0] |
|
|
fmri_data = data_file[key] |
|
|
fmri_data = fmri_data.astype(np.float32) |
|
|
except Exception as e: |
|
|
print(f"Error loading file {file_path}: {e}") |
|
|
return None |
|
|
|
|
|
total_time_frames = fmri_data.shape[-1] |
|
|
if total_time_frames > self.crop_length: |
|
|
start_idx = np.random.randint(0, total_time_frames - self.crop_length + 1) |
|
|
end_idx = start_idx + self.crop_length |
|
|
cropped_data = fmri_data[..., start_idx:end_idx] |
|
|
else: |
|
|
cropped_data = fmri_data[..., :self.crop_length] |
|
|
|
|
|
data_tensor = torch.from_numpy(cropped_data) |
|
|
|
|
|
data_tensor = data_tensor.permute(3, 0, 1, 2) |
|
|
|
|
|
return data_tensor |
|
|
|
