Upload 3 files

configs/data_configs.py

@@ -0,0 +1,41 @@
+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['celeba_test'],
+		'test_target_root': dataset_paths['celeba_test'],
+	},
+	'cars_encode': {
+		'transforms': transforms_config.CarsEncodeTransforms,
+		'train_source_root': dataset_paths['cars_train'],
+		'train_target_root': dataset_paths['cars_train'],
+		'test_source_root': dataset_paths['cars_test'],
+		'test_target_root': dataset_paths['cars_test'],
+	},
+	'horse_encode': {
+		'transforms': transforms_config.EncodeTransforms,
+		'train_source_root': dataset_paths['horse_train'],
+		'train_target_root': dataset_paths['horse_train'],
+		'test_source_root': dataset_paths['horse_test'],
+		'test_target_root': dataset_paths['horse_test'],
+	},
+	'church_encode': {
+		'transforms': transforms_config.EncodeTransforms,
+		'train_source_root': dataset_paths['church_train'],
+		'train_target_root': dataset_paths['church_train'],
+		'test_source_root': dataset_paths['church_test'],
+		'test_target_root': dataset_paths['church_test'],
+	},
+	'cats_encode': {
+		'transforms': transforms_config.EncodeTransforms,
+		'train_source_root': dataset_paths['cats_train'],
+		'train_target_root': dataset_paths['cats_train'],
+		'test_source_root': dataset_paths['cats_test'],
+		'test_target_root': dataset_paths['cats_test'],
+	}
+}

configs/paths_config.py

@@ -0,0 +1,28 @@
+dataset_paths = {
+	#  Face Datasets (In the paper: FFHQ - train, CelebAHQ - test)
+	'ffhq': '',
+	'celeba_test': '',
+
+	#  Cars Dataset (In the paper: Stanford cars)
+	'cars_train': '',
+	'cars_test': '',
+
+	#  Horse Dataset (In the paper: LSUN Horse)
+	'horse_train': '',
+	'horse_test': '',
+
+	#  Church Dataset (In the paper: LSUN Church)
+	'church_train': '',
+	'church_test': '',
+
+	#  Cats Dataset (In the paper: LSUN Cat)
+	'cats_train': '',
+	'cats_test': ''
+}
+
+model_paths = {
+	'stylegan_ffhq': 'pretrained_models/stylegan2-ffhq-config-f.pt',
+	'ir_se50': 'pretrained_models/model_ir_se50.pth',
+	'shape_predictor': 'pretrained_models/shape_predictor_68_face_landmarks.dat',
+	'moco': 'pretrained_models/moco_v2_800ep_pretrain.pth'
+}

configs/transforms_config.py

@@ -0,0 +1,62 @@
+from abc import abstractmethod
+import torchvision.transforms as transforms
+
+
+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((256, 256)),
+				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((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 CarsEncodeTransforms(TransformsConfig):
+
+	def __init__(self, opts):
+		super(CarsEncodeTransforms, self).__init__(opts)
+
+	def get_transforms(self):
+		transforms_dict = {
+			'transform_gt_train': transforms.Compose([
+				transforms.Resize((192, 256)),
+				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((192, 256)),
+				transforms.ToTensor(),
+				transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+			'transform_inference': transforms.Compose([
+				transforms.Resize((192, 256)),
+				transforms.ToTensor(),
+				transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+		}
+		return transforms_dict