import os
import torch
import torchvision
import tarfile
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from torchvision.datasets.utils import download_url
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import torchvision.transforms as tt
from torch.utils.data import random_split
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
%matplotlib inline

data_dir = '../input/face-mask-12k-images-dataset/Face Mask Dataset'
print(os.listdir(data_dir))
classes = os.listdir(data_dir + "/Train")
print(classes)

train_ds = ImageFolder(data_dir+'/Train')
val_ds = ImageFolder(data_dir+'/Validation')
test_ds = ImageFolder(data_dir+'/Test')

print(len(train_ds))
print(len(val_ds))
print(len(test_ds))

train_ds.classes