Created
July 17, 2019 00:08
-
-
Save Hanrui-Wang/0bfec986c7031a2405a3fb8d6559ffc2 to your computer and use it in GitHub Desktop.
how to create a dataset in the pytorch
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| class FaceLandmarksDataset(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_csv(csv_file) | |
| self.root_dir = root_dir | |
| self.transform = transform | |
| def __len__(self): | |
| return len(self.landmarks_frame) | |
| def __getitem__(self, idx): | |
| img_name = os.path.join(self.root_dir, | |
| self.landmarks_frame.iloc[idx, 0]) | |
| image = io.imread(img_name) | |
| landmarks = self.landmarks_frame.iloc[idx, 1:] | |
| landmarks = np.array([landmarks]) | |
| landmarks = landmarks.astype('float').reshape(-1, 2) | |
| sample = {'image': image, 'landmarks': landmarks} | |
| if self.transform: | |
| sample = self.transform(sample) | |
| return sample |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment