Open andife opened 3 years ago
The following dataset.py is working for me:
import torch.utils.data
import torchvision
from torchvision import dataset
import cv2
cv2.setNumThreads(0)
dataset = datasets.ImageFolder("/Project/")
images = dataset.imgs
class SearchDataset(torch.utils.data.Dataset):
def __init__(self, images_filepaths=images, transform=None):
self.images_filepaths = images_filepaths
self.transform = transform
def __len__(self):
return len(self.images_filepaths)
def __getitem__(self, index):
path, label = self.images_filepaths[index]
image = cv2.imread(path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
if self.transform is not None:
transformed = self.transform(image=image)
image = transformed["image"]
return image, label
Images should be all of the same size. If they are not, you can resize them by adding to __getitem__
:
image = cv2.resize(image, (height, width))
Hello, very interesting package.
Can someone provide a code snipped for a "Add implementation for len and getitem methods in dataset.py."
My dataset is a classification problem, and the files are structured like
Project/group1/bmp Project/group2/bmp
Thank you
Andreas