|
from torch.utils import data |
|
from PIL import Image |
|
import os |
|
|
|
|
|
class Dataset(data.Dataset): |
|
'Characterizes a dataset for PyTorch' |
|
|
|
def __init__(self, path, transform=None): |
|
'Initialization' |
|
self.file_names = self.get_filenames(path) |
|
self.transform = transform |
|
|
|
def __len__(self): |
|
'Denotes the total number of samples' |
|
return len(self.file_names) |
|
|
|
def __getitem__(self, index): |
|
'Generates one sample of data' |
|
img = Image.open(self.file_names[index]).convert('RGB') |
|
|
|
if self.transform is not None: |
|
img = self.transform(img) |
|
return img |
|
|
|
def get_filenames(self, data_path): |
|
images = [] |
|
for path, subdirs, files in os.walk(data_path): |
|
for name in files: |
|
if name.rfind('jpg') != -1 or name.rfind('png') != -1: |
|
filename = os.path.join(path, name) |
|
if os.path.isfile(filename): |
|
images.append(filename) |
|
return images |
|
|
