Here is my custom DATASET
In each sample, I have the image data and other data (float number)
I hope I can do the normalization for image
, tr_angles
and pls
.
Thanks a lot.
class PLDataset(Dataset):
"""Face Landmarks dataset."""
def __init__(self, csv_file, root_dir, transform=None):
"""
Args:
csv_file (string): Path to the csv file with annotations.
root_dir (string): Directory with all the images.
transform (callable, optional): Optional transform to be applied
on a sample.
"""
self.landmarks_frame = pd.read_excel(csv_file)
self.root_dir = root_dir
self.transform = transform
def __len__(self):
return len(self.landmarks_frame)
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = idx.tolist()
img_names = os.path.join(self.root_dir,
self.landmarks_frame.iloc[idx, 0])
image = io.imread(img_names)
pls = self.landmarks_frame.iloc[idx, 5:9] # not [:, 5:9]
tr_angles = self.landmarks_frame.iloc[idx, 9:14]
positions = self.landmarks_frame.iloc[idx, 3:5]
image = image.astype('float') # .values
pls = pls.values.astype('float')
tr_angles = tr_angles.values.astype('float')
positions = positions.values.astype('float')
sample = {'image': image,
'pls': pls,
'tr_angles': tr_angles,
'positions': positions}
if self.transform:
sample = self.transform(sample)
return sample
Now I just have my own ToTensor()
function,
class ToTensor(object):
"""Convert ndarrays in sample to Tensors."""
def __call__(self, sample):
image, pls, tr_angles, positions = \
sample['image'], sample['pls'], sample['tr_angles'], sample['positions']
image = image[np.newaxis, :, :]
# swap color axis because
# numpy image: H x W x C
# torch image: C X H X W
# image = image.transpose((2, 0, 1)) # single channel
return {'image': torch.from_numpy(image),
'pls': torch.from_numpy(pls),
'tr_angles': torch.from_numpy(tr_angles),
'positions': torch.from_numpy(positions)
}