Skip to content

Commit

Permalink
add testdataset class
Browse files Browse the repository at this point in the history
  • Loading branch information
snoop2head committed Sep 3, 2021
1 parent fe72103 commit fec0f74
Showing 1 changed file with 18 additions and 0 deletions.
18 changes: 18 additions & 0 deletions dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,3 +108,21 @@ class GenderLabels:

class AgeGroup:
map_label = lambda x: 0 if int(x) < 30 else 1 if int(x) < 60 else 2


class TestDataset(Dataset):
def __init__(self, img_paths, transform, device):
self.img_paths = img_paths
self.device = device
self.transform = transform

def __getitem__(self, index):
image = Image.open(self.img_paths[index])

if self.transform:
image = self.transform(image=np.array(image))['image']

return image.to(self.device)

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

0 comments on commit fec0f74

Please sign in to comment.