fix lsun data bug

Author: xperzy

gan/celeba_dataset.py gan/Styleformer/celeba_dataset.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 """
-CelebA Dataset related classes and methods 
+CelebA Dataset related classes and methods
 Currently only support for GAN
 """
 
@@ -44,17 +44,17 @@ def __len__(self):
         return len(self.img_path_list)
 
     def __getitem__(self, index):
-        data = Image.open(self.img_path_list[index]).convert('RGB')
+        img = Image.open(self.img_path_list[index]).convert('RGB')
         if self.transform is not None:
-            data = self.transform(data)
+            img = self.transform(img)
         label = 0
-        return data, label
-
-if __name__ == "__main__":
-    dataset = CelebADataset(file_folder='./celeba/img_align_celeba')
-    for idx, (data, label) in enumerate(dataset):
-        print(idx)
-        print(data.size)
-        print('-----')
-        if idx == 10:
-            break 
+        return img, label
+
+#if __name__ == "__main__":
+#    dataset = CelebADataset(file_folder='./celeba/img_align_celeba')
+#    for idx, (data, label) in enumerate(dataset):
+#        print(idx)
+#        print(data.size)
+#        print('-----')
+#        if idx == 10:
+#            break

gan/Styleformer/config.py

@@ -17,7 +17,6 @@
 Configuration for data, model archtecture, and training, etc.
 Config can be set by .yaml file or by argparser(limited usage)
 
-
 """
 
 import os
@@ -37,6 +36,8 @@
 _C.DATA.CHANNEL = 3 # input image channel
 _C.DATA.CROP_PCT = 1.0 # input image scale ratio, scale is applied before centercrop in eval mode
 _C.DATA.NUM_WORKERS = 2 # number of data loading threads
+_C.DATA.MAX_REAL_NUM = None # number of images used in the dataset (real images)
+_C.DATA.MAX_GEN_NUM = None # number of images used in the generator (fake images)
 
 # model settings
 _C.MODEL = CN()
@@ -73,7 +74,6 @@
 _C.TRAIN.WARMUP_START_LR = 0.0
 _C.TRAIN.END_LR = 0.0
 _C.TRAIN.GRAD_CLIP = 1.0
-_C.TRAIN.ACCUM_ITER = 2
 
 _C.TRAIN.LR_SCHEDULER = CN()
 _C.TRAIN.LR_SCHEDULER.NAME = 'warmupcosine'
@@ -161,7 +161,9 @@ def update_config(config, args):
     return config
 
 
-def get_config():
-    """Return a clone config"""
+def get_config(cfg_file=None):
+    """Return a clone of config or load from yaml file"""
     config = _C.clone()
+    if cfg_file:
+        _update_config_from_file(config, cfg_file)
     return config

gan/Styleformer/configs/styleformer_celeba.yaml

@@ -1,8 +1,11 @@
 DATA:
     IMAGE_SIZE: 64
+    MAX_GEN_NUM: 50000
+    MAX_REAL_NUM: None
 MODEL:
     TYPE: Styleformer
     NAME: Styleformer_Linformer
+    NUM_CLASSES: 10177
     GEN:
         RESOLUTION: 8
         NUM_LAYERS: [1,2,1,1]

gan/Styleformer/configs/styleformer_cifar10.yaml

@@ -1,8 +1,11 @@
 DATA:
     IMAGE_SIZE: 32
+    MAX_GEN_NUM: 50000
+    MAX_REAL_NUM: None
 MODEL:
     TYPE: Styleformer
     NAME: Styleformer_Large
+    NUM_CLASSES: 10
     GEN:
         RESOLUTION: 8
         NUM_LAYERS: [1,3,3]

gan/Styleformer/configs/styleformer_lsun.yaml

@@ -1,8 +1,11 @@
 DATA:
     IMAGE_SIZE: 128
+    MAX_GEN_NUM: 50000
+    MAX_REAL_NUM: None
 MODEL:
     TYPE: Styleformer
     NAME: Styleformer_Linformer
+    NUM_CLASSES: 1
     GEN:
         RESOLUTION: 8
         NUM_LAYERS: [1,2,1,1,1]

gan/Styleformer/configs/styleformer_stl10.yaml

@@ -1,8 +1,11 @@
 DATA:
     IMAGE_SIZE: 48
+    MAX_GEN_NUM: 50000
+    MAX_REAL_NUM: None
 MODEL:
     TYPE: Styleformer
     NAME: Styleformer_Medium
+    NUM_CLASSES: 1 # unlabeled data, all class 0
     GEN:
         RESOLUTION: 12
         NUM_LAYERS: [1,3,3]

gan/Styleformer/datasets.py

@@ -19,14 +19,15 @@
 
 import os
 import math
-import sys
-import numpy as np
-from paddle.io import Dataset, DataLoader, DistributedBatchSampler
-from paddle.vision import transforms, datasets, image_load
-sys.path.append('../')
-from stl10_dataset import *
-from lsun_church_dataset import *
-from celeba_dataset import *
+from paddle.io import Dataset
+from paddle.io import DataLoader
+from paddle.io import DistributedBatchSampler
+from paddle.vision import transforms
+from paddle.vision import datasets
+from paddle.vision import image_load
+from stl10_dataset import STL10Dataset
+from lsun_church_dataset import LSUNchurchDataset
+from celeba_dataset import CelebADataset
 
 class ImageNet2012Dataset(Dataset):
     """Build ImageNet2012 dataset
@@ -140,7 +141,8 @@ def get_dataset(config, mode='train'):
                                    mode=mode,
                                    transform=get_train_transforms(config))
         else:
-            mode = 'test'
+            #mode = 'test'
+            mode = 'unlabeled'
             dataset = STL10Dataset(file_folder=config.DATA.DATA_PATH,
                                    mode=mode,
                                    transform=get_val_transforms(config))
@@ -156,10 +158,10 @@ def get_dataset(config, mode='train'):
     elif config.DATA.DATASET == "celeba":
         if mode == 'train':
             dataset = CelebADataset(file_folder=config.DATA.DATA_PATH,
-                                        transform=get_train_transforms(config))
+                                    transform=get_train_transforms(config))
         else:
             dataset = CelebADataset(file_folder=config.DATA.DATA_PATH,
-                                        transform=get_val_transforms(config))
+                                    transform=get_val_transforms(config))
     elif config.DATA.DATASET == "imagenet2012":
         if mode == 'train':
             dataset = ImageNet2012Dataset(config.DATA.DATA_PATH,
@@ -171,7 +173,7 @@ def get_dataset(config, mode='train'):
                                           transform=get_val_transforms(config))
     else:
         raise NotImplementedError(
-            "[{config.DATA.DATASET}] Only cifar10, cifar100, imagenet2012 are supported now")
+            "Only support cifar10, cifar100, imagenet2012, celeba, stl10, lsun")
     return dataset
 
 

gan/Styleformer/discriminator.py

@@ -1,4 +1,4 @@
-#   Copyright (c) 2021 PPViT Authors. All Rights Reserved.
+# Copyright (c) 2021 PPViT Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -167,4 +167,4 @@ def forward(self, input):
         out = out.reshape((batch, -1))
         out = self.final_linear(out)
 
-        return out
+        return out

gan/Styleformer/generate.py

@@ -1,4 +1,4 @@
-#   Copyright (c) 2021 PPViT Authors. All Rights Reserved.
+# Copyright (c) 2021 PPViT Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,47 +12,68 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Generate images using trained models"""
 import argparse
-import numpy
-import paddle
+import os
 from PIL import Image
+import paddle
 from generator import Generator
-from config import *
+from config import get_config
+from config import update_config
 
 
 def main():
-    # get default config
-    parser = argparse.ArgumentParser('')
+    """ generate sample images using pretrained model
+    The following args are required:
+        -cfg: str, path of yaml model config file
+        -pretrained: str, path of the pretrained model (ends with .pdparams)
+        -num_out_images: int, the num of output images to be saved in file
+        -out_folder: str, output folder path.
+    """
+    paddle.set_device('gpu')
+    # get config
+    parser = argparse.ArgumentParser('Generate samples images')
     parser.add_argument('-cfg', type=str, default='./configs/styleformer_cifar10.yaml')
-    parser.add_argument('-dataset', type=str, default="cifar10")
+    parser.add_argument('-pretrained', type=str, default='./lsun.pdparams')
+    parser.add_argument('-num_out_images', type=int, default=16)
+    parser.add_argument('-out_folder', type=str, default='./out_images_lsun')
+
+    parser.add_argument('-dataset', type=str, default=None)
     parser.add_argument('-batch_size', type=int, default=None)
     parser.add_argument('-image_size', type=int, default=None)
     parser.add_argument('-ngpus', type=int, default=None)
-    parser.add_argument('-data_path', type=str, default='/dataset/cifar10/')
+    parser.add_argument('-data_path', type=str, default=None)
     parser.add_argument('-eval', action="store_true")
-    parser.add_argument('-pretrained', type=str, default=None)
+
     args = parser.parse_args()
     config = get_config()
     config = update_config(config, args)
-
-    paddle.set_device('cpu')
+    # get model
+    print(f'----- Creating model...')
     paddle_model = Generator(config)
     paddle_model.eval()
-
-    pre=paddle.load(r'./cifar10.pdparams')
-    paddle_model.load_dict(pre)
-
-    x = paddle.randn([32, 512])
-    x_paddle = paddle.to_tensor(x)
-    out_paddle = paddle_model(x_paddle, c=paddle.randint(0, 10, [32]))
-
-    gen_imgs=paddle.multiply(out_paddle,paddle.to_tensor(127.5))
-    gen_imgs=paddle.clip(paddle.add(gen_imgs,paddle.to_tensor(127.5)).transpose((0,2,3,1)),
-             min=0.0,max=255.0).astype('uint8').cpu().numpy()
-
-    for i in range(len(gen_imgs)):
-        im = Image.fromarray(gen_imgs[i], 'RGB')
-        im.save("./image/"+str(i)+".png")
+    # load model weights
+    print(f'----- Loading model form {config.MODEL.PRETRAINED}...')
+    model_state_dict = paddle.load(config.MODEL.PRETRAINED)
+    paddle_model.load_dict(model_state_dict)
+    # get random input tensor
+    x_paddle = paddle.randn([args.num_out_images, paddle_model.z_dim])
+    # inference
+    print(f'----- Inferencing...')
+    out_paddle = paddle_model(
+        z=x_paddle, c=paddle.randint(0, config.MODEL.NUM_CLASSES, [args.num_out_images]))
+    # post processing to obtain image
+    print('----- Postprocessing')
+    gen_imgs = (out_paddle * 127.5 + 128).clip(0, 255)
+    gen_imgs = gen_imgs.transpose((0, 2, 3, 1)).astype('uint8')
+    gen_imgs = gen_imgs.cpu().numpy()
+    # save images to file
+    os.makedirs(args.out_folder, exist_ok=True)
+    print(f'----- Saving images to {args.out_folder}')
+    for i, gen_img in enumerate(gen_imgs):
+        img = Image.fromarray(gen_img, 'RGB')
+        out_path = os.path.join(args.out_folder, str(i) + '.png')
+        img.save(out_path)
 
 
 if __name__ == "__main__":