Reputation: 31
I am doing image classification with PyTorch. I have a separate Images folder and train and test csv file with images ids and labels . I don’t have any an idea about how to combine those images and ID and converting into tensors.
Upvotes: 3
Views: 8386
Reputation: 5412
You can create custom dataset class by inherting pytorch's torch.utils.data.Dataset.
The assumption for the following custom dataset class is
filename | label |
---|---|
4325.jpg | cat |
2345.jpg | dog |
images folder
.class CustomDataset(torch.utils.data.Dataset):
def __init__(self, csv_path, images_folder, transform = None):
self.df = pd.read_csv(csv_path)
self.images_folder = images_folder
self.transform = transform
self.class2index = {"cat":0, "dog":1}
def __len__(self):
return len(self.df)
def __getitem__(self, index):
filename = self.df[index, "FILENAME"]
label = self.class2index[self.df[index, "LABEL"]]
image = PIL.Image.open(os.path.join(self.images_folder, filename))
if self.transform is not None:
image = self.transform(image)
return image, label
Now you can use this class to load the training and test dataset using both csv file and image folder.
train_dataset = CustomDataset("path - to - train.csv", "path - to - images - folder" )
test_dataset = CustomDataset("path - to - test.csv", "path - to - images - folder" )
image, label = train_dataset[0]
Upvotes: 8