I tried the folloiwng based off of https://jhui.github.io/2018/02/09/PyTorch-Data-loading-preprocess_torchvision/ 156 but got an error. Code as below:

class Roof_Dataset(Dataset):
“”“Roof dataset.”""
def __init__(self, csv_file, root_dir, transform):
    """
    Args:
        csv_file (string): Path to the csv file with annotations.
        root_dir (string): Directory with all the images.
        transform (callable, optional): Optional transform to be applied
            on a sample.
    """
    
    self.roofs_frame = pd.read_csv(csv_file)
    self.root_dir = root_dir
    self.transform = transform

def __len__(self):
    return len(self.landmarks_frame)

def __getitem__(self, idx):
    img_name = Path(self.root_dir,
                            self.roof_frame.iloc[idx, 1])
    image = io.imread(img_name)
    landmarks = self.landmarks_frame.iloc[idx, 5]
    sample = {'image': image, 'landmarks': landmarks}

    if self.transform:
        sample = self.transform(sample)

    return sample


I get  the following;


NameError                                 Traceback (most recent call last)
 in 
----> 1 class Roof_Dataset(Dataset):
2     “”“Roof dataset.”""
3
4     def init(self, csv_file, root_dir, transform):
5         “”"
NameError: name ‘Dataset’ is not defined

I’m also not sure, how I would have the data loader randomly cut off 20% for a test set. Am I headed in the wrong direction?