import os

import os.path

import numpy as np

import random

import h5py

import torch

import glob

import torch.utils.data as udata

class MRI_Dataset(udata.Dataset):

    def __init__(self, x_path, y_path):

        super(MRI_Dataset, self).__init__()

        self.x_path = x_path;

        self.y_path = y_path;

        h5f_x = h5py.File(self.x_path, 'r')

        h5f_y = h5py.File(self.y_path, 'r')

        self.keys = list(h5f_x.keys())

        random.shuffle(self.keys)

        h5f_x.close()

        h5f_y.close()

    def __len__(self):

        return len(self.keys)

    def __getitem__(self, index):

        h5f_x = h5py.File(self.x_path, 'r')

        h5f_y = h5py.File(self.y_path, 'r')

        key = self.keys[index]

        data_x = np.array(h5f_x[key])

        data_y = np.array(h5f_y[key])

        h5f_x.close();

        h5f_y.close();

        return (torch.Tensor(data_x), torch.Tensor(data_y),key)