61228f3243eaee39cba31a9abd2afc6ead9612a8,implementations/discogan/datasets.py,ImageDataset,__getitem__,#ImageDataset#Any#,15

Before Change

    def __getitem__(self, index):

        img_pair = self.transform(Image.open(self.files[index % len(self.files)]))
        _, h, w = img_pair.shape
        half_w = int(w/2)

        item_A = img_pair[:, :, :half_w]
        item_B = img_pair[:, :, half_w:]

        return {"A": item_A, "B": item_B}

    def __len__(self):

After Change

        self.files = sorted(glob.glob(os.path.join(root, mode) + "/*.*"))

    def __getitem__(self, index):

        img = Image.open(self.files[index % len(self.files)])
        w, h = img.size
        img_A = img.crop((0, 0, w/2, h))
        img_B = img.crop((w/2, 0, w, h))

        if np.random.random() < 0.5:
            img_A = Image.fromarray(np.array(img_A)[:, ::-1, :], "RGB")
            img_B = Image.fromarray(np.array(img_B)[:, ::-1, :], "RGB")

        img_A = self.transform(img_A)
        img_B = self.transform(img_B)

        return {"A": img_A, "B": img_B}

    def __len__(self):

Instances