pytorch_paintings_dataset.py

class SquarePad:
    def __call__(self, image):
        w, h = image.size
        max_wh = np.max([w, h])
        hp = int((max_wh - w) / 2)
        vp = int((max_wh - h) / 2)
        padding = (hp, vp, hp, vp)
        return torchvision.transforms.functional.pad(image, padding, 0, 'constant')
    
class ArtPeriodDataSet(Dataset):
    def __init__(self, dataframe, img_dir, transform=None, target_transform=None):
        self.dataframe = dataframe.copy()
        self.img_dir = img_dir
        # eliminate non-existing files
        print("Eliminate non-existing files")
        for idx, row in tqdm(dataframe.iterrows(), total=len(dataframe)):
            path = f'{img_dir}/{row["ID"]}.jpg'
            if not os.path.exists(path):
                self.dataframe.drop(idx, inplace=True)
                print("Dropping", path)
        
        # encode brands
        self.label_encoder = LabelEncoder()
        labels_encoded = self.label_encoder.fit_transform(self.dataframe["period"])
        self.dataframe["period_encoded"] = labels_encoded
        self.classes = self.label_encoder.classes_
        
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

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

    def __getitem__(self, idx):        
        img_path = f'{self.img_dir}/{self.dataframe.iloc[idx]["ID"]}.jpg'
        image = PIL.Image.open(img_path).convert('RGB')
        label = self.dataframe.iloc[idx]["period_encoded"]
        
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label