(self,
data_root,
size=None,
repeats=100,
interpolation="bicubic",
flip_p=0.5,
set="train",
placeholder_token="*",
per_image_tokens=False,
center_crop=False,
mixing_prob=0.25,
coarse_class_text=None,
)
| 131 | |
| 132 | class PersonalizedBase(Dataset): |
| 133 | def __init__(self, |
| 134 | data_root, |
| 135 | size=None, |
| 136 | repeats=100, |
| 137 | interpolation="bicubic", |
| 138 | flip_p=0.5, |
| 139 | set="train", |
| 140 | placeholder_token="*", |
| 141 | per_image_tokens=False, |
| 142 | center_crop=False, |
| 143 | mixing_prob=0.25, |
| 144 | coarse_class_text=None, |
| 145 | ): |
| 146 | |
| 147 | self.data_root = data_root |
| 148 | |
| 149 | self.image_paths = [os.path.join(self.data_root, file_path) for file_path in os.listdir(self.data_root)] |
| 150 | |
| 151 | # self._length = len(self.image_paths) |
| 152 | self.num_images = len(self.image_paths) |
| 153 | self._length = self.num_images |
| 154 | |
| 155 | self.placeholder_token = placeholder_token |
| 156 | |
| 157 | self.per_image_tokens = per_image_tokens |
| 158 | self.center_crop = center_crop |
| 159 | self.mixing_prob = mixing_prob |
| 160 | |
| 161 | self.coarse_class_text = coarse_class_text |
| 162 | |
| 163 | if per_image_tokens: |
| 164 | assert self.num_images < len(per_img_token_list), f"Can't use per-image tokens when the training set contains more than {len(per_img_token_list)} tokens. To enable larger sets, add more tokens to 'per_img_token_list'." |
| 165 | |
| 166 | if set == "train": |
| 167 | self._length = self.num_images * repeats |
| 168 | |
| 169 | self.size = size |
| 170 | self.interpolation = {"linear": PIL.Image.LINEAR, |
| 171 | "bilinear": PIL.Image.BILINEAR, |
| 172 | "bicubic": PIL.Image.BICUBIC, |
| 173 | "lanczos": PIL.Image.LANCZOS, |
| 174 | }[interpolation] |
| 175 | self.flip = transforms.RandomHorizontalFlip(p=flip_p) |
| 176 | |
| 177 | def __len__(self): |
| 178 | return self._length |
nothing calls this directly
no outgoing calls
no test coverage detected