Modified the data_folder to be more generic

Hey @ageorge,

It's a good practice to open a merge request each time you intend to work on branch. Doing this way we can discuss changes in an easier way.

Concerning what you did, are you sure that it still works with Olegs stuff ?

Thanks

bob/learn/pytorch/datasets/data_folder.py

@@ -12,9 +12,8 @@ import torch.utils.data as data
 import os
 
 import random
-random.seed( a = 7 )
 
-import PIL
+random.seed( a = 7 )
 
 import numpy as np
 
@@ -65,11 +64,11 @@ def get_file_names_and_labels(files, data_folder, extension = ".hdf5", hldi_type
 
             if f.attack_type is None:
 
-                label = 0
+                label = 1
 
             else:
 
-                label = 1
+                label = 0
 
             file_name = os.path.join(data_folder, f.path + extension)
 
@@ -102,7 +101,7 @@ class DataFolder(data.Dataset):
         A directory containing the training data.
 
     transform : object
-        A function/transform that  takes in a PIL image, and returns a
+        A function/transform that  takes in a numpy.ndarray image, and returns a
         transformed version. E.g, ``transforms.RandomCrop``. Default: None.
 
     extension : str
@@ -155,7 +154,7 @@ class DataFolder(data.Dataset):
             A directory containing the training data.
 
         transform : object
-            A function/transform that  takes in a PIL image, and returns a
+            A custom function/transform (torchvision.transforms.Compose) that takes in a numpy.ndarray image , and returns a
             transformed version. E.g, ``transforms.RandomCrop``. Default: None.
 
         extension : str
@@ -238,9 +237,9 @@ class DataFolder(data.Dataset):
         Returns
         -------
 
-        pil_img : Tensor or PIL Image
-            If ``self.transform`` is defined the output is the torch.Tensor,
-            otherwise the output is an instance of the PIL.Image.Image class.
+        np_img : Tensor image
+            If ``self.transform`` is defined as a custom function, the output is the torch.Tensor,
+            otherwise the last transform in the transforms.Compose should be transforms.ToTensor().
 
         target : int
             Index of the class.
@@ -254,20 +253,21 @@ class DataFolder(data.Dataset):
 
         if isinstance(self.transform, transforms.Compose): # if an instance of torchvision composed transformation
 
-            if len(img_array.shape) == 3: # for color images
+            if len(img_array.shape) == 3: # for color or multi-channel images
 
                 img_array_tr = np.swapaxes(img_array, 1, 2)
                 img_array_tr = np.swapaxes(img_array_tr, 0, 2)
 
-                pil_img = PIL.Image.fromarray( img_array_tr ) # convert to PIL from array of size (H x W x 3)
+                np_img =img_array_tr.copy()  # np_img is numpy.ndarray of shape HxWxC
 
             else: # for gray-scale images
 
-                pil_img = PIL.Image.fromarray( img_array, 'L' ) # convert to PIL from array of size (H x W)
+                np_img=np.expand_dims(img_array_tr,2) # np_img is numpy.ndarray of size HxWx1
+
 
             if self.transform is not None:
 
-                pil_img = self.transform(pil_img)
+                np_img = self.transform(np_img) # after this transformation np_img should be a tensor
 
         else: # if custom transformation function is given
 
@@ -275,7 +275,7 @@ class DataFolder(data.Dataset):
 
             return torch.Tensor(img_array_transformed).unsqueeze(0), target # convert array to Tensor, also return target
 
-        return pil_img, target
+        return np_img, target
 
 
     #==========================================================================