Method __getitem__

src/finetune_data.py:190–325

(self, i)

            return self._length1

    def __getitem__(self, i):
        example = {}
        mask_background = False
        example["mask_background_image"] = np.ones((self.size // 8, self.size // 8))
        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)) ]
            
            if self.mask_background:
                mask_background_image = Image.open(self.mask_paths[i % self._length1])
                mask_background = True
        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_prev = image
        image = self.flip(image)

        

        


        if mask_background:
            mbi = np.array(mask_background_image).astype(np.uint8)
            crop = min(mbi.shape[0], mbi.shape[1])
            h, w, = mbi.shape[0], mbi.shape[1]

            mbi = mbi[(h - crop) // 2:(h + crop) // 2,
                    (w - crop) // 2:(w + crop) // 2]
            mask_background_image = Image.fromarray(mbi)            
            diff = ImageChops.difference(image_prev, image)
            if diff.getbbox():
                mask_background_image = self.mask_flip(mask_background_image)
            if self.size is not None:
                mask_background_image = mask_background_image.resize((self.size // 8, self.size // 8), resample=Image.NEAREST)
            
            mask_background_image = np.array(mask_background_image).astype(np.uint8)
            mask_background_image = (mask_background_image / 255).astype(np.float32)