I’ve created a custom dataset that allows me to read in images from a single folder and add the labels separately:
class SingleFolderImage(Dataset):
def __init__(self, base_path, label_dict, transform=None):
self.base_path = base_path
self.image_paths = [str(f) for f in Path(self.base_path).glob('*')]
self.transform = transform
self.label_dict = label_dict
def __getitem__(self, index):
image_path = self.image_paths[index]
image_name = re.split(r'/|\.', image_path)[-2]
x = cv2.imread(image_path, -1)
ratio = np.amax(x) / 256
x = (x / ratio).astype('uint8')
if self.transform is not None:
x = self.transform(x)
if image_name in list(self.label_dict.keys()):
y = self.label_dict[image_name]
else:
y = {'boxes': None, 'labels': None}
return x, y
def __len__(self):
return len(self.image_paths)
I call it as following:
train_dataset = SingleFolderImage(image_path, parallel_labels, transform = data_transforms)
train_dataloader = DataLoader(full_train_dataset, batch_size=100, shuffle = True, num_workers=2)
However, when I try to read some examples I get the error in the title, without any useful way to isolate where the issue is. This didn’t happen before I added the if statement that handles images that aren’t in label_dict (but of course that leads to other issues), but I really don’t see why that would affect anything, after all I’m returning a y with the same type in both cases.
Any help would be most welcome…