Keras to PyTorch - DatasetFolder/ImageFolder/Custom?

Huhu,
I am currently trying to reconstruct my auto encoder in PyTorch, but I have some problems loading the data or maybe problems understanding. can someone help me on the jumps ?

TF/Keras:

from tensorflow.keras.preprocessing.image import load_img, img_to_array
from sklearn.model_selection import train_test_split

train_images_path = '../denoising/train'
train_cleaned_path = '../denoising/train_cleaned'

train_images = sorted(os.listdir(train_images_path))
train_cleaned = sorted(os.listdir(train_cleaned_path))

x_train = []
y_train = []

for image in train_images:
    img_path = os.path.join(train_images_path, image)
    img = load_img(img_path, color_mode = 'grayscale', target_size = (420, 540))
    img = img_to_array(img).astype('float32')/255
    x_train.append(img)
for image in train_cleaned:
    img_path = os.path.join(train_cleaned_path, image)
    img = load_img(img_path, color_mode = 'grayscale', target_size = (420, 540))
    img = img_to_array(img).astype('float32')/255
    y_train.append(img)
     
x = np.array(x_train)
y = np.array(y_train)

x_train, x_valid, y_train, y_valid = train_test_split(x, y, test_size = 0.25)

PyTorch:

from tensorflow.keras.preprocessing.image import load_img, img_to_array
from sklearn.model_selection import train_test_split

train_images_path = '../denoising/train'
train_cleaned_path = '../denoising/train_cleaned'

train_images = sorted(os.listdir(train_images_path))
train_cleaned = sorted(os.listdir(train_cleaned_path))

x_train = []
y_train = []

for image in train_images:
    img_path = os.path.join(train_images_path, image)
    img = load_img(img_path, color_mode = 'grayscale', target_size = (420, 540))
    img = img_to_array(img).astype('float32')/255
    x_train.append(img)
for image in train_cleaned:
    img_path = os.path.join(train_cleaned_path, image)
    img = load_img(img_path, color_mode = 'grayscale', target_size = (420, 540))
    img = img_to_array(img).astype('float32')/255
    y_train.append(img)
     
x = np.array(x_train)
y = np.array(y_train)

x_train, x_valid, y_train, y_valid = train_test_split(x, y, test_size = 0.25)

transform = transforms.Compose([
     transforms.ToPILImage(),
     transforms.ToTensor()
])

class ImageData(Dataset):
    def __init__(self, images, labels=None, transforms=None):
        self.X = images
        self.y = labels
        self.transforms = transforms
         
    def __len__(self):
        return (len(self.X))
    
    def __getitem__(self, i):
        data = self.X[i][:]
        
        if self.transforms:
            data = self.transforms(data)
            
        if self.y is not None:
            labels = self.y[i][:]
            labels = self.transforms(labels)
            return (data, labels)
        else:
            return data
        
train_data = ImageData(x_train, y_train, transform)
test_data = ImageData(x_valid, None, transform)
train_loader = DataLoader(train_data, batch_size=BATCH_SIZE)
test_loader = DataLoader(test_data, batch_size=BATCH_SIZE)
  1. Isn’t that easier in “pure” PyTorch without OpenCV or else ?

  2. Assuming I want to assign a label to every image how could I do that with the ImageFolder?
    Keras:

for f in os.listdir(best_doc_dir):
        img_file = os.path.join(best_doc_dir, f)
        try:
            img = load_img(img_file, color_mode = 'grayscale', target_size = (IMG_HEIGHT, IMG_WIDTH))
            img = img_to_array(img).astype('float32')/255
            x[cnt] = img
            y[cnt] = BEST_CLASS_IDX
            cnt += 1
        except:
            print("best document image %s cannot be read!" % f)

    # Dropping not readable image idxs
    x = x[:cnt]
    y = y[:cnt]

    np.save(os.path.join(FILE_DIR, "x.npy"), x)
    np.save(os.path.join(FILE_DIR, "y.npy"), y)

Thanks for your help :slight_smile:

You PyTorch code looks generally alright, although you could lazily load the data using a custom Dataset as described in this tutorial.

Inside this dataset you can also assign your custom labels in the __getitem__ method. Let me know, if you get stuck. :wink: