| 35 | |
| 36 | |
| 37 | class CelebA(torch.utils.data.Dataset): |
| 38 | def __init__(self, root, transform): |
| 39 | |
| 40 | self.root = root |
| 41 | self.transform = transform |
| 42 | self.base_folder = "celeba" |
| 43 | |
| 44 | fn = partial(os.path.join, self.root, self.base_folder) |
| 45 | splits = pd.read_csv(fn("list_eval_partition.txt"), |
| 46 | delim_whitespace=True, header=None, index_col=0) |
| 47 | # identity = pandas.read_csv(fn("identity_CelebA.txt"), delim_whitespace=True, header=None, index_col=0) |
| 48 | # bbox = pandas.read_csv(fn("list_bbox_celeba.txt"), delim_whitespace=True, header=1, index_col=0) |
| 49 | # landmarks_align = pandas.read_csv(fn("list_landmarks_align_celeba.txt"), delim_whitespace=True, header=1) |
| 50 | self.attr = pd.read_csv( |
| 51 | fn("list_attr_celeba.txt"), delim_whitespace=True, header=1) |
| 52 | # change label -1 to label 0 |
| 53 | self.attr = self.attr.replace(-1, 0) |
| 54 | |
| 55 | # mask = slice(None) |
| 56 | # filename: ['000001.jpg' '000002.jpg' '000003.jpg' ... '202597.jpg' '202598.jpg' '202599.jpg'] |
| 57 | self.filename = splits.index.values |
| 58 | |
| 59 | self.attr_names = list(self.attr.columns) |
| 60 | |
| 61 | def __getitem__(self, index): |
| 62 | X = Image.open(os.path.join(self.root, self.base_folder, |
| 63 | "img_align_celeba", self.filename[index])) |
| 64 | # X = jpeg.JPEG(os.path.join(self.root, self.base_folder, "img_align_celeba", self.filename[index])).decode() |
| 65 | # X = self.image_list[index] |
| 66 | target = {k: self.attr.loc[self.filename[index]][k] |
| 67 | for k in self.attr_names} |
| 68 | target["Mouth_Smile"] = target["Mouth_Slightly_Open"] * \ |
| 69 | 2 + target["Smiling"] |
| 70 | target["Young_Smile"] = target["Young"] * 2 + target["Smiling"] |
| 71 | target["Attractive_Smile"] = target["Attractive"] * \ |
| 72 | 2 + target["Smiling"] |
| 73 | |
| 74 | if self.transform is not None: |
| 75 | X = self.transform(X) |
| 76 | |
| 77 | return X, target |
| 78 | |
| 79 | def __len__(self): |
| 80 | return len(self.attr) |
| 81 | |
| 82 | |
| 83 | if __name__ == "__main__": |
no outgoing calls
no test coverage detected