# Define custom dataset
class Predict_Set(Dataset):
    def __init__(self, img_folder, transform):
        self.img_folder = img_folder
        self.transform = transform
        self.img_lst = os.listdir(self.img_folder)

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

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_folder, self.img_lst[idx])
        img = Image.open(img_path)
        img = img.crop((50, 60, 750, 460))  #Size: 700x400
        img.load()
        img_tensor = self.transform(img)
        return img_tensor, self.img_lst[idx]