Upload 5 files

configs/data_configs.py

@@ -0,0 +1,48 @@
+from configs import transforms_config
+from configs.paths_config import dataset_paths
+
+
+DATASETS = {
+    'ffhq_encode': {
+        'transforms': transforms_config.EncodeTransforms,
+        'train_source_root': dataset_paths['ffhq'],
+        'train_target_root': dataset_paths['ffhq'],
+        'test_source_root': dataset_paths['ffhq_test'],
+        'test_target_root': dataset_paths['ffhq_test'],
+    },
+    'ffhq_sketch_to_face': {
+        'transforms': transforms_config.SketchToImageTransforms,
+        'train_source_root': dataset_paths['ffhq_train_sketch'],
+        'train_target_root': dataset_paths['ffhq'],
+        'test_source_root': dataset_paths['ffhq_test_sketch'],
+        'test_target_root': dataset_paths['ffhq_test'],
+    },
+    'ffhq_seg_to_face': {
+        'transforms': transforms_config.SegToImageTransforms,
+        'train_source_root': dataset_paths['ffhq_train_segmentation'],
+        'train_target_root': dataset_paths['ffhq'],
+        'test_source_root': dataset_paths['ffhq_test_segmentation'],
+        'test_target_root': dataset_paths['ffhq_test'],
+    },
+    'ffhq_super_resolution': {
+        'transforms': transforms_config.SuperResTransforms,
+        'train_source_root': dataset_paths['ffhq'],
+        'train_target_root': dataset_paths['ffhq1280'],
+        'test_source_root': dataset_paths['ffhq_test'],
+        'test_target_root': dataset_paths['ffhq1280_test'],
+    },
+    'toonify': {
+        'transforms': transforms_config.ToonifyTransforms,
+        'train_source_root': dataset_paths['toonify_in'],
+        'train_target_root': dataset_paths['toonify_out'],
+        'test_source_root': dataset_paths['toonify_test_in'],
+        'test_target_root': dataset_paths['toonify_test_out'],
+    },   
+    'ffhq_edit': {
+        'transforms': transforms_config.EditingTransforms,
+        'train_source_root': dataset_paths['ffhq'],
+        'train_target_root': dataset_paths['ffhq'],
+        'test_source_root': dataset_paths['ffhq_test'],
+        'test_target_root': dataset_paths['ffhq_test'],
+    },
+}

configs/dataset_config.yml

@@ -0,0 +1,60 @@
+# dataset and data loader settings
+datasets:
+  train:
+    name: FFHQ
+    type: FFHQDegradationDataset
+    # dataroot_gt: datasets/ffhq/ffhq_512.lmdb
+    dataroot_gt: ../../../../share/shuaiyang/ffhq/realign1280x1280test/
+    io_backend:
+      # type: lmdb
+      type: disk
+
+    use_hflip: true
+    mean: [0.5, 0.5, 0.5]
+    std: [0.5, 0.5, 0.5]
+    out_size: 1280
+    scale: 4
+
+    blur_kernel_size: 41
+    kernel_list: ['iso', 'aniso']
+    kernel_prob: [0.5, 0.5]
+    blur_sigma: [0.1, 10]
+    downsample_range: [4, 40]
+    noise_range: [0, 20]
+    jpeg_range: [60, 100]
+
+    # color jitter and gray
+    #color_jitter_prob: 0.3
+    #color_jitter_shift: 20
+    #color_jitter_pt_prob: 0.3
+    #gray_prob: 0.01
+
+    # If you do not want colorization, please set
+    color_jitter_prob: ~
+    color_jitter_pt_prob: ~
+    gray_prob: 0.01
+    gt_gray: True
+
+    crop_components: true
+    component_path: ./pretrained_models/FFHQ_eye_mouth_landmarks_512.pth
+    eye_enlarge_ratio: 1.4
+
+    # data loader
+    use_shuffle: true
+    num_worker_per_gpu: 6
+    batch_size_per_gpu: 4
+    dataset_enlarge_ratio: 1
+    prefetch_mode: ~
+
+  val:
+    # Please modify accordingly to use your own validation
+    # Or comment the val block if do not need validation during training
+    name: validation
+    type: PairedImageDataset
+    dataroot_lq: datasets/faces/validation/input
+    dataroot_gt: datasets/faces/validation/reference
+    io_backend:
+      type: disk
+    mean: [0.5, 0.5, 0.5]
+    std: [0.5, 0.5, 0.5]
+    scale: 1

configs/paths_config.py

@@ -0,0 +1,25 @@
+dataset_paths = {
+    'ffhq': 'data/train/ffhq/realign320x320/',
+    'ffhq_test': 'data/train/ffhq/realign320x320test/',
+    'ffhq1280': 'data/train/ffhq/realign1280x1280/',
+    'ffhq1280_test': 'data/train/ffhq/realign1280x1280test/',    
+    'ffhq_train_sketch': 'data/train/ffhq/realign640x640sketch/',
+    'ffhq_test_sketch': 'data/train/ffhq/realign640x640sketchtest/',
+    'ffhq_train_segmentation': 'data/train/ffhq/realign320x320mask/',
+    'ffhq_test_segmentation': 'data/train/ffhq/realign320x320masktest/',
+    'toonify_in': 'data/train/pixar/trainA/',
+    'toonify_out': 'data/train/pixar/trainB/',
+    'toonify_test_in': 'data/train/pixar/testA/',
+    'toonify_test_out': 'data/train/testB/',
+}
+
+model_paths = {
+    'stylegan_ffhq': 'pretrained_models/stylegan2-ffhq-config-f.pt',
+    'ir_se50': 'pretrained_models/model_ir_se50.pth',
+    'circular_face': 'pretrained_models/CurricularFace_Backbone.pth',
+    'mtcnn_pnet': 'pretrained_models/mtcnn/pnet.npy',
+    'mtcnn_rnet': 'pretrained_models/mtcnn/rnet.npy',
+    'mtcnn_onet': 'pretrained_models/mtcnn/onet.npy',
+    'shape_predictor': 'shape_predictor_68_face_landmarks.dat',
+    'moco': 'pretrained_models/moco_v2_800ep_pretrain.pth.tar'
+}

configs/transforms_config.py

@@ -0,0 +1,242 @@
+from abc import abstractmethod
+import torchvision.transforms as transforms
+from datasets import augmentations
+
+
+class TransformsConfig(object):
+
+    def __init__(self, opts):
+        self.opts = opts
+
+    @abstractmethod
+    def get_transforms(self):
+        pass
+
+
+class EncodeTransforms(TransformsConfig):
+
+    def __init__(self, opts):
+        super(EncodeTransforms, self).__init__(opts)
+
+    def get_transforms(self):
+        transforms_dict = {
+            'transform_gt_train': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.RandomHorizontalFlip(0.5),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_source': None,
+            'transform_test': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_inference': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+        }
+        return transforms_dict
+
+
+class FrontalizationTransforms(TransformsConfig):
+
+    def __init__(self, opts):
+        super(FrontalizationTransforms, self).__init__(opts)
+
+    def get_transforms(self):
+        transforms_dict = {
+            'transform_gt_train': transforms.Compose([
+                transforms.Resize((256, 256)),
+                transforms.RandomHorizontalFlip(0.5),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_source': transforms.Compose([
+                transforms.Resize((256, 256)),
+                transforms.RandomHorizontalFlip(0.5),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_test': transforms.Compose([
+                transforms.Resize((256, 256)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_inference': transforms.Compose([
+                transforms.Resize((256, 256)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+        }
+        return transforms_dict
+
+
+class SketchToImageTransforms(TransformsConfig):
+
+    def __init__(self, opts):
+        super(SketchToImageTransforms, self).__init__(opts)
+
+    def get_transforms(self):
+        transforms_dict = {
+            'transform_gt_train': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_source': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor()]),
+            'transform_test': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_inference': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor()]),
+        }
+        return transforms_dict
+
+
+class SegToImageTransforms(TransformsConfig):
+
+    def __init__(self, opts):
+        super(SegToImageTransforms, self).__init__(opts)
+
+    def get_transforms(self):
+        transforms_dict = {
+            'transform_gt_train': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_source': transforms.Compose([
+                transforms.Resize((320, 320)),
+                augmentations.ToOneHot(self.opts.label_nc),
+                transforms.ToTensor()]),
+            'transform_test': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_inference': transforms.Compose([
+                transforms.Resize((320, 320)),
+                augmentations.ToOneHot(self.opts.label_nc),
+                transforms.ToTensor()])
+        }
+        return transforms_dict
+
+
+class SuperResTransforms(TransformsConfig):
+
+    def __init__(self, opts):
+        super(SuperResTransforms, self).__init__(opts)
+
+    def get_transforms(self):
+        if self.opts.resize_factors is None:
+            self.opts.resize_factors = '1,2,4,8,16,32'
+        factors = [int(f) for f in self.opts.resize_factors.split(",")]
+        print("Performing down-sampling with factors: {}".format(factors))
+        transforms_dict = {
+            'transform_gt_train': transforms.Compose([
+                transforms.Resize((1280, 1280)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_source': transforms.Compose([
+                transforms.Resize((320, 320)),
+                augmentations.BilinearResize(factors=factors),
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_test': transforms.Compose([
+                transforms.Resize((1280, 1280)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_inference': transforms.Compose([
+                transforms.Resize((320, 320)),
+                augmentations.BilinearResize(factors=factors),
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+        }
+        return transforms_dict
+
+    
+class SuperResTransforms_320(TransformsConfig):
+
+    def __init__(self, opts):
+        super(SuperResTransforms_320, self).__init__(opts)
+
+    def get_transforms(self):
+        if self.opts.resize_factors is None:
+            self.opts.resize_factors = '1,2,4,8,16,32'
+        factors = [int(f) for f in self.opts.resize_factors.split(",")]
+        print("Performing down-sampling with factors: {}".format(factors))
+        transforms_dict = {
+            'transform_gt_train': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_source': transforms.Compose([
+                transforms.Resize((320, 320)),
+                augmentations.BilinearResize(factors=factors),
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_test': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_inference': transforms.Compose([
+                transforms.Resize((320, 320)),
+                augmentations.BilinearResize(factors=factors),
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+        }
+        return transforms_dict
+    
+
+class ToonifyTransforms(TransformsConfig):
+
+    def __init__(self, opts):
+        super(ToonifyTransforms, self).__init__(opts)
+
+    def get_transforms(self):
+        transforms_dict = {
+            'transform_gt_train': transforms.Compose([
+                transforms.Resize((1024, 1024)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_source': transforms.Compose([
+                transforms.Resize((256, 256)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_test': transforms.Compose([
+                transforms.Resize((1024, 1024)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_inference': transforms.Compose([
+                transforms.Resize((256, 256)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+        }
+        return transforms_dict
+    
+class EditingTransforms(TransformsConfig):
+
+    def __init__(self, opts):
+        super(EditingTransforms, self).__init__(opts)
+
+    def get_transforms(self):
+        transforms_dict = {
+            'transform_gt_train': transforms.Compose([
+                transforms.Resize((1280, 1280)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_source': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_test': transforms.Compose([
+                transforms.Resize((1280, 1280)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+            'transform_inference': transforms.Compose([
+                transforms.Resize((320, 320)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+        }
+        return transforms_dict