Fix bugs in jpeg data loading

data/__init__.py

+from .datasets import *
\ No newline at end of file

data/datasets.py

 import torch
 from torchvision import datasets, transforms
-from jpeg_dataset import JpegDataset
+from .jpeg_dataset import JpegDataset
 
 
 def spatial_data(batch_size, root, name, dataset, transform, shuffle_train):
@@ -64,3 +64,31 @@ def cifar10_jpeg(batch_size, root, shuffle_train=True):
 def cifar100_jpeg(batch_size, root, shuffle_train=True):
     directory = '{}/{}'.format(root, 'CIFAR100')
     return jpeg_data(batch_size, directory, shuffle_train)
+
+
+spatial_dataset_map = {
+    'MNIST': mnist_spatial,
+    'CIFAR10': cifar10_spatial,
+    'CIFAR100': cifar100_spatial
+}
+
+jpeg_dataset_map = {
+    'MNIST': mnist_jpeg,
+    'CIFAR10': cifar10_jpeg,
+    'CIFAR100': cifar100_jpeg
+}
+
+dataset_info = {
+    'MNIST': {
+        'channels': 1,
+        'classes': 10
+    },
+    'CIFAR10': {
+        'channels': 3,
+        'classes': 10
+    },
+    'CIFAR100': {
+        'channels': 3,
+        'classes': 100,
+    }
+}

data/jpeg_dataset.py

 import torch
-from jpeg_readwrite import load_jpeg
+from .jpeg_readwrite import load_jpeg
 
 
 class JpegDataset(torch.utils.data.Dataset):
@@ -10,4 +10,4 @@ class JpegDataset(torch.utils.data.Dataset):
         return self.data.shape[0]
 
     def __getitem__(self, idx):
-        return self.dataset[idx], self.labels[idx]
+        return self.data[idx], self.labels[idx]

data/jpeg_readwrite.py

@@ -13,10 +13,9 @@ def write_jpeg(data, labels, directory, fname):
 def load_jpeg(directory, fname):
     path = '{}/{}'.format(directory, fname)
 
-    with open(path) as f:
+    with open(path, 'rb') as f:
         dl = np.load(f)
+        data = dl['data']
+        labels = dl['labels']
 
-    data = dl['data']
-    label = dl['label']
-
-    return torch.Tensor(data), torch.Tensor(label)
+    return torch.Tensor(data), torch.Tensor(labels)