Face Landmarks dataset.
| 118 | # |
| 119 | |
| 120 | class FaceLandmarksDataset(Dataset): |
| 121 | """Face Landmarks dataset.""" |
| 122 | |
| 123 | def __init__(self, csv_file, root_dir, transform=None): |
| 124 | """ |
| 125 | Arguments: |
| 126 | csv_file (string): Path to the csv file with annotations. |
| 127 | root_dir (string): Directory with all the images. |
| 128 | transform (callable, optional): Optional transform to be applied |
| 129 | on a sample. |
| 130 | """ |
| 131 | self.landmarks_frame = pd.read_csv(csv_file) |
| 132 | self.root_dir = root_dir |
| 133 | self.transform = transform |
| 134 | |
| 135 | def __len__(self): |
| 136 | return len(self.landmarks_frame) |
| 137 | |
| 138 | def __getitem__(self, idx): |
| 139 | if torch.is_tensor(idx): |
| 140 | idx = idx.tolist() |
| 141 | |
| 142 | img_name = os.path.join(self.root_dir, |
| 143 | self.landmarks_frame.iloc[idx, 0]) |
| 144 | image = io.imread(img_name) |
| 145 | landmarks = self.landmarks_frame.iloc[idx, 1:] |
| 146 | landmarks = np.array([landmarks], dtype=float).reshape(-1, 2) |
| 147 | sample = {'image': image, 'landmarks': landmarks} |
| 148 | |
| 149 | if self.transform: |
| 150 | sample = self.transform(sample) |
| 151 | |
| 152 | return sample |
| 153 | |
| 154 | |
| 155 | ###################################################################### |
no outgoing calls
no test coverage detected