(self, index)
| 334 | return self._length |
| 335 | |
| 336 | def __getitem__(self, index): |
| 337 | example = {} |
| 338 | instance_image, instance_prompt = self.instance_images_path[index % self.num_instance_images] |
| 339 | instance_image = Image.open(instance_image) |
| 340 | if not instance_image.mode == "RGB": |
| 341 | instance_image = instance_image.convert("RGB") |
| 342 | instance_image = self.flip(instance_image) |
| 343 | |
| 344 | ############################################################################## |
| 345 | #### apply resize augmentation and create a valid image region mask ########## |
| 346 | ############################################################################## |
| 347 | if np.random.randint(0, 3) < 2: |
| 348 | random_scale = np.random.randint(self.size // 3, self.size+1) |
| 349 | else: |
| 350 | random_scale = np.random.randint(int(1.2*self.size), int(1.4*self.size)) |
| 351 | |
| 352 | if random_scale % 2 == 1: |
| 353 | random_scale += 1 |
| 354 | |
| 355 | if random_scale < 0.6*self.size: |
| 356 | add_to_caption = np.random.choice(["a far away ", "very small "]) |
| 357 | instance_prompt = add_to_caption + instance_prompt |
| 358 | cx = np.random.randint(random_scale // 2, self.size - random_scale // 2 + 1) |
| 359 | cy = np.random.randint(random_scale // 2, self.size - random_scale // 2 + 1) |
| 360 | instance_image1 = preprocess(instance_image, random_scale, self.interpolation) |
| 361 | instance_image = np.zeros((self.size, self.size, 3), dtype=np.float32) |
| 362 | instance_image[cx - random_scale // 2: cx + random_scale // 2, cy - random_scale // 2: cy + random_scale // 2, :] = instance_image1 |
| 363 | |
| 364 | mask = np.zeros((self.size // 8, self.size // 8)) |
| 365 | mask[(cx - random_scale // 2) // 8 + 1: (cx + random_scale // 2) // 8 - 1, (cy - random_scale // 2) // 8 + 1: (cy + random_scale // 2) // 8 - 1] = 1. |
| 366 | elif random_scale > self.size: |
| 367 | add_to_caption = np.random.choice(["zoomed in ", "close up "]) |
| 368 | instance_prompt = add_to_caption + instance_prompt |
| 369 | cx = np.random.randint(self.size // 2, random_scale - self.size // 2 + 1) |
| 370 | cy = np.random.randint(self.size // 2, random_scale - self.size // 2 + 1) |
| 371 | |
| 372 | instance_image = preprocess(instance_image, random_scale, self.interpolation) |
| 373 | instance_image = instance_image[cx - self.size // 2: cx + self.size // 2, cy - self.size // 2: cy + self.size // 2, :] |
| 374 | mask = np.ones((self.size // 8, self.size // 8)) |
| 375 | else: |
| 376 | instance_image = preprocess(instance_image, self.size, self.interpolation) |
| 377 | mask = np.ones((self.size // 8, self.size // 8)) |
| 378 | ######################################################################## |
| 379 | |
| 380 | example["instance_images"] = torch.from_numpy(instance_image).permute(2, 0, 1) |
| 381 | example["mask"] = torch.from_numpy(mask) |
| 382 | example["instance_prompt_ids"] = self.tokenizer( |
| 383 | instance_prompt, |
| 384 | truncation=True, |
| 385 | padding="max_length", |
| 386 | max_length=self.tokenizer.model_max_length, |
| 387 | return_tensors="pt", |
| 388 | ).input_ids |
| 389 | |
| 390 | if self.with_prior_preservation: |
| 391 | class_image, class_prompt = self.class_images_path[index % self.num_class_images] |
| 392 | class_image = Image.open(class_image) |
| 393 | if not class_image.mode == "RGB": |
nothing calls this directly
no test coverage detected