ADNI_dataset.py

import csv
import numpy as np
import torch
from torch.utils.data.dataset import Dataset
import os
from torchvision import transforms
from skimage.transform import resize
from nilearn import surface
import nibabel as nib
from ipdb import set_trace

class ADNIdataset(Dataset):
	def __init__(self, root='../ADNI', augmentation=False, img_size=64, normalization=True):
		self.normalization=normalization
		self.img_size=img_size
		self.root = root
		self.basis = 'FreeSurfer_Cross-Sectional_Processing_brainmask'
		self.augmentation = augmentation
		f = open('CN_list.csv','r')
		rdr = csv.reader(f)
		name = []
		labels = []
		date  = []
		for line in rdr:
			[month,day,year] = line[9].split('/')
			month = month.zfill(2)
			date.append(year+'-'+month+'-'+day)
			name.append(line[1])

		name = np.asarray(name)
		date = np.asarray(date)

		self.name =name
		self.date =date
	def __len__(self):
		return len(self.name)

	def __getitem__(self, index):
		path = os.path.join(self.root,self.name[index],self.basis)
		files = os.listdir(path)
		for file in files:
			if file[:10] == self.date[index]:
				rname = file
		aname = os.listdir(os.path.join(path,rname))[0]
		path = os.path.join(path,rname,aname,'mri')
		img = nib.load(os.path.join(path,'brainmask.mgz'))
		img = np.swapaxes(img.get_data(),1,2)
		img = np.flip(img,1)
		img = np.flip(img,2)
		img = resize(img, (self.img_size,self.img_size,self.img_size), mode='constant')
		if self.augmentation:
			random_n = torch.rand(1)
			random_i = 0.3*torch.rand(1)[0]+0.7
			if random_n[0] > 0.5:
				img = np.flip(img,0)
			img = img*random_i.data.cpu().numpy()     
		imageout = img
		if self.normalization: 
			imageout = torch.from_numpy(img).float().view(1,self.img_size,self.img_size,self.img_size)
			imageout = imageout*2-1

		return imageout