#Create Dataset
path = "/content/drive/MyDrive/Colab Notebooks/LIDAR/SPAD_NYU/SPAD_Counts"
path_short = "/content/drive/MyDrive/Colab Notebooks/LIDAR/SPAD_NYU/SPAD_Counts_Short"

class SPADSet(Dataset):
  def __init__(self,path,transform=None,is_train=False):
    super(SPADSet, self).__init__()
    self.path = path
    self.fileList = list()
    for file in os.listdir(path):
      self.fileList.append(path+"/"+file)
    self.transform = transform
    self.is_train = is_train

  def __len__(self):
    return len(self.fileList)

  def __getitem__(self, idx):
    fileName = self.fileList[idx]
    spadpath = fileName

    #Load MATLAB file, then get return tuple of spad and depth.
    mat = scipy.io.loadmat(spadpath)
    spad = mat['spad']
    spad = spad.astype(int)
    spad = torch.tensor(spad)
    #spad = torch.unsqueeze(spad,dim=0)
    spad = spad.float()
    spad = torch.reshape(spad,(128,64,64))
    depth = mat['depth']
    depth = torch.tensor(depth)
    depth = torch.unsqueeze(depth,dim=0)
    depth = depth.float()

    return spad, depth

dataset = SPADSet(path)
print(len(dataset))

train_data,valid_data = torch.utils.data.random_split(dataset, [350, 24])

train_dataloader = DataLoader(train_data, batch_size=4, shuffle=True)
test_dataloader = DataLoader(valid_data, batch_size=4, shuffle=True)

#CREATE ITERATOR
train_iterator = iter(train_dataloader)
test_iterator = iter(test_dataloader)