fad73753f836de5cdd08d8e23d26c1170d19b386,src/dataset_unpair.py,dataset_unpair,__getitem__,#dataset_unpair#,21

Before Change

  def __getitem__(self, index):
    crop_img = self.load_img(self.images[index], flip=self.flip, random_crop=True)
    raw_data = crop_img.transpose((2, 0, 1))
    data = ((torch.FloatTensor(raw_data) / 255.0) - 0.5) * 2
    return data

  def load_img(self, img_name, flip=False, random_crop=False):
    if self.input_dim == 3:
      img = cv2.cvtColor(cv2.imread(img_name), cv2.COLOR_BGR2RGB)

After Change

  def __getitem__(self, index):
    if self.dataset_size == self.A_size:
      data_A = self.load_img(self.A[index], self.input_dim_A)
      data_B = self.load_img(self.B[random.randint(0, self.B_size - 1)], self.input_dim_B)
    else:
      data_A = self.load_img(self.A[random.randint(0, self.A_size - 1)], self.input_dim_A)
      data_B = self.load_img(self.B[index], self.input_dim_B)
    return data_A, data_B

  def load_img(self, img_name, input_dim):
    img = Image.open(img_name).convert("RGB")
    img = self.transforms(img)

Instances