Method __getitem__

src/finetune_data.py:182–267

(self, i)

            return self._length1

    def __getitem__(self, i):
        example = {}

        if i > self._length2 or self._length2 == 0:
            image = Image.open(self.labels["relative_file_path1_"][i % self._length1])
            if isinstance(self.caption, str):
                example["caption"] = np.random.choice(self.templates_small).format(self.caption)
            else:
                example["caption"] = self.caption[i % min(self._length1, len(self.caption)) ]
        else:
            image = Image.open(self.labels["relative_file_path2_"][i % self._length2])
            if isinstance(self.reg_caption, str):
                example["caption"] = np.random.choice(self.templates_small).format(self.reg_caption)
            else:
                example["caption"] = self.reg_caption[i % self._length2]

        if not image.mode == "RGB":
            image = image.convert("RGB")

        # default to score-sde preprocessing
        img = np.array(image).astype(np.uint8)
        crop = min(img.shape[0], img.shape[1])
        h, w, = img.shape[0], img.shape[1]

        img = img[(h - crop) // 2:(h + crop) // 2,
                  (w - crop) // 2:(w + crop) // 2]

        image = Image.fromarray(img)
        image = self.flip(image)

        if i > self._length2 or self._length2 == 0:
            if self.aug:
                if np.random.randint(0, 3) < 2:
                    random_scale = np.random.randint(self.size // 3, self.size+1)
                else:
                    random_scale = np.random.randint(int(1.2*self.size), int(1.4*self.size))

                if random_scale % 2 == 1:
                    random_scale += 1
            else:
                random_scale = self.size

            if random_scale < 0.6*self.size:
                add_to_caption = np.random.choice(["a far away ", "very small "])
                example["caption"] = add_to_caption + example["caption"]
                cx = np.random.randint(random_scale // 2, self.size - random_scale // 2 + 1)
                cy = np.random.randint(random_scale // 2, self.size - random_scale // 2 + 1)

                image = image.resize((random_scale, random_scale), resample=self.interpolation)
                image = np.array(image).astype(np.uint8)
                image = (image / 127.5 - 1.0).astype(np.float32)

                input_image1 = np.zeros((self.size, self.size, 3), dtype=np.float32)
                input_image1[cx - random_scale // 2: cx + random_scale // 2, cy - random_scale // 2: cy + random_scale // 2, :] = image

                mask = np.zeros((self.size // 8, self.size // 8))
                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.