").append(n.parseHTML(a)).find(d):a)}).always(c&&function(a,b){g.each(function(){c.apply(this,f||[a.responseText,b,a])})}),this},n.each(["ajaxStart","ajaxStop","ajaxComplete","ajaxError","ajaxSuccess","ajaxSend"],function(a,b){n.fn[b]=function(a){return this.on(b,a)}}),n.expr.filters.animated=function(a){return n.grep(n.timers,function(b){return a===b.elem}).length};function mc(a){return n.isWindow(a)?a:9===a.nodeType?a.defaultView||a.parentWindow:!1}n.offset={setOffset:function(a,b,c){var d,e,f,g,h,i,j,k=n.css(a,"position"),l=n(a),m={};"static"===k&&(a.style.position="relative"),h=l.offset(),f=n.css(a,"top"),i=n.css(a,"left"),j=("absolute"===k||"fixed"===k)&&n.inArray("auto",[f,i])>-1,j?(d=l.position(),g=d.top,e=d.left):(g=parseFloat(f)||0,e=parseFloat(i)||0),n.isFunction(b)&&(b=b.call(a,c,n.extend({},h))),null!=b.top&&(m.top=b.top-h.top+g),null!=b.left&&(m.left=b.left-h.left+e),"using"in b?b.using.call(a,m):l.css(m)}},n.fn.extend({offset:function(a){if(arguments.length)return void 0===a?this:this.each(function(b){n.offset.setOffset(this,a,b)});var b,c,d={top:0,left:0},e=this[0],f=e&&e.ownerDocument;if(f)return b=f.documentElement,n.contains(b,e)?("undefined"!=typeof e.getBoundingClientRect&&(d=e.getBoundingClientRect()),c=mc(f),{top:d.top+(c.pageYOffset||b.scrollTop)-(b.clientTop||0),left:d.left+(c.pageXOffset||b.scrollLeft)-(b.clientLeft||0)}):d},position:function(){if(this[0]){var a,b,c={top:0,left:0},d=this[0];return"fixed"===n.css(d,"position")?b=d.getBoundingClientRect():(a=this.offsetParent(),b=this.offset(),n.nodeName(a[0],"html")||(c=a.offset()),c.top+=n.css(a[0],"borderTopWidth",!0),c.left+=n.css(a[0],"borderLeftWidth",!0)),{top:b.top-c.top-n.css(d,"marginTop",!0),left:b.left-c.left-n.css(d,"marginLeft",!0)}}},offsetParent:function(){return this.map(function(){var a=this.offsetParent;while(a&&!n.nodeName(a,"html")&&"static"===n.css(a,"position"))a=a.offsetParent;return a||Qa})}}),n.each({scrollLeft:"pageXOffset",scrollTop:"pageYOffset"},function(a,b){var c=/Y/.test(b);n.fn[a]=function(d){return Y(this,function(a,d,e){var f=mc(a);return void 0===e?f?b in f?f[b]:f.document.documentElement[d]:a[d]:void(f?f.scrollTo(c?n(f).scrollLeft():e,c?e:n(f).scrollTop()):a[d]=e)},a,d,arguments.length,null)}}),n.each(["top","left"],function(a,b){n.cssHooks[b]=Ua(l.pixelPosition,function(a,c){return c?(c=Sa(a,b),Oa.test(c)?n(a).position()[b]+"px":c):void 0})}),n.each({Height:"height",Width:"width"},function(a,b){n.each({
+padding:"inner"+a,content:b,"":"outer"+a},function(c,d){n.fn[d]=function(d,e){var f=arguments.length&&(c||"boolean"!=typeof d),g=c||(d===!0||e===!0?"margin":"border");return Y(this,function(b,c,d){var e;return n.isWindow(b)?b.document.documentElement["client"+a]:9===b.nodeType?(e=b.documentElement,Math.max(b.body["scroll"+a],e["scroll"+a],b.body["offset"+a],e["offset"+a],e["client"+a])):void 0===d?n.css(b,c,g):n.style(b,c,d,g)},b,f?d:void 0,f,null)}})}),n.fn.extend({bind:function(a,b,c){return this.on(a,null,b,c)},unbind:function(a,b){return this.off(a,null,b)},delegate:function(a,b,c,d){return this.on(b,a,c,d)},undelegate:function(a,b,c){return 1===arguments.length?this.off(a,"**"):this.off(b,a||"**",c)}}),n.fn.size=function(){return this.length},n.fn.andSelf=n.fn.addBack,"function"==typeof define&&define.amd&&define("jquery",[],function(){return n});var nc=a.jQuery,oc=a.$;return n.noConflict=function(b){return a.$===n&&(a.$=oc),b&&a.jQuery===n&&(a.jQuery=nc),n},b||(a.jQuery=a.$=n),n});
diff --git a/SimSwap/docs/js/which-image.js b/SimSwap/docs/js/which-image.js
new file mode 100644
index 0000000000000000000000000000000000000000..867bf202b687b5e0e403f040e0688e44a27ea0f3
--- /dev/null
+++ b/SimSwap/docs/js/which-image.js
@@ -0,0 +1,50 @@
+/*
+ * @FilePath: \SimSwap\docs\js\which-image.js
+ * @Author: Ziang Liu
+ * @Date: 2021-07-03 16:34:56
+ * @LastEditors: AceSix
+ * @LastEditTime: 2021-07-20 00:46:27
+ * Copyright (C) 2021 SJTU. All rights reserved.
+ */
+
+
+
+function select_source(number) {
+ var items = ['anni', 'chenglong', 'zhoujielun', 'zhuyin'];
+ var item_id = items[number];
+
+ for (i = 0; i < 4; i++) {
+ if (number == i) {
+ document.getElementById(items[i]).style.borderWidth = '5px';
+ document.getElementById(items[i]).style.borderColor = 'red';
+ document.getElementById(items[i]).style.borderStyle = 'outset';
+ } else {
+ document.getElementById(items[i]).style.border = 'none';
+ }
+ }
+ document.getElementById('jiroujinlun').src = './img/' + item_id + '.webp';
+
+}
+
+function select_source2(number) {
+ var items = ['Iron_man', 'wuyifan'];
+ var item_id = items[number];
+
+ for (i = 0; i < 2; i++) {
+ if (number == i) {
+ document.getElementById(items[i]).style.borderWidth = '5px';
+ document.getElementById(items[i]).style.borderColor = 'red';
+ document.getElementById(items[i]).style.borderStyle = 'outset';
+ } else {
+ document.getElementById(items[i]).style.border = 'none';
+ }
+ }
+ if (item_id=='Iron_man'){
+ document.getElementById('mama').src = './img/mama_mask_short.webp';
+ }
+ else{
+ document.getElementById('mama').src = './img/mama_mask_wuyifan_short.webp';
+
+ }
+
+}
\ No newline at end of file
diff --git a/SimSwap/download-weights.sh b/SimSwap/download-weights.sh
new file mode 100644
index 0000000000000000000000000000000000000000..54e6ea5c68575b7f2a79b20386b53f8d19e2c62c
--- /dev/null
+++ b/SimSwap/download-weights.sh
@@ -0,0 +1,9 @@
+#!/bin/sh
+wget -P ./arcface_model https://github.com/neuralchen/SimSwap/releases/download/1.0/arcface_checkpoint.tar
+wget https://github.com/neuralchen/SimSwap/releases/download/1.0/checkpoints.zip
+unzip ./checkpoints.zip -d ./checkpoints
+rm checkpoints.zip
+wget --no-check-certificate "https://sh23tw.dm.files.1drv.com/y4mmGiIkNVigkSwOKDcV3nwMJulRGhbtHdkheehR5TArc52UjudUYNXAEvKCii2O5LAmzGCGK6IfleocxuDeoKxDZkNzDRSt4ZUlEt8GlSOpCXAFEkBwaZimtWGDRbpIGpb_pz9Nq5jATBQpezBS6G_UtspWTkgrXHHxhviV2nWy8APPx134zOZrUIbkSF6xnsqzs3uZ_SEX_m9Rey0ykpx9w" -O antelope.zip
+mkdir -p insightface_func/models
+unzip ./antelope.zip -d ./insightface_func/models/
+rm antelope.zip
diff --git a/SimSwap/insightface_func/__init__.py b/SimSwap/insightface_func/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/SimSwap/insightface_func/__pycache__/__init__.cpython-310.pyc b/SimSwap/insightface_func/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..91b9d050cf0f31ecae2553a84c021776ec64dac9
Binary files /dev/null and b/SimSwap/insightface_func/__pycache__/__init__.cpython-310.pyc differ
diff --git a/SimSwap/insightface_func/__pycache__/face_detect_crop_multi.cpython-310.pyc b/SimSwap/insightface_func/__pycache__/face_detect_crop_multi.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..843bb99f19c2dc0ac49955c9e4468e6b2c3bc02a
Binary files /dev/null and b/SimSwap/insightface_func/__pycache__/face_detect_crop_multi.cpython-310.pyc differ
diff --git a/SimSwap/insightface_func/__pycache__/face_detect_crop_single.cpython-310.pyc b/SimSwap/insightface_func/__pycache__/face_detect_crop_single.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bac1369711345f4453defbc5f98ad3d860fac0a1
Binary files /dev/null and b/SimSwap/insightface_func/__pycache__/face_detect_crop_single.cpython-310.pyc differ
diff --git a/SimSwap/insightface_func/face_detect_crop_multi.py b/SimSwap/insightface_func/face_detect_crop_multi.py
new file mode 100644
index 0000000000000000000000000000000000000000..d75c36b58baeaa8c49e15dc9c3f0986076eb5fdd
--- /dev/null
+++ b/SimSwap/insightface_func/face_detect_crop_multi.py
@@ -0,0 +1,100 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 17:03:58
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-24 16:45:41
+Description:
+'''
+from __future__ import division
+import collections
+import numpy as np
+import glob
+import os
+import os.path as osp
+import cv2
+from insightface.model_zoo import model_zoo
+from insightface_func.utils import face_align_ffhqandnewarc as face_align
+
+__all__ = ['Face_detect_crop', 'Face']
+
+Face = collections.namedtuple('Face', [
+ 'bbox', 'kps', 'det_score', 'embedding', 'gender', 'age',
+ 'embedding_norm', 'normed_embedding',
+ 'landmark'
+])
+
+Face.__new__.__defaults__ = (None, ) * len(Face._fields)
+
+
+class Face_detect_crop:
+ def __init__(self, name, root='~/.insightface_func/models'):
+ self.models = {}
+ root = os.path.expanduser(root)
+ onnx_files = glob.glob(osp.join(root, name, '*.onnx'))
+ onnx_files = sorted(onnx_files)
+ for onnx_file in onnx_files:
+ if onnx_file.find('_selfgen_')>0:
+ #print('ignore:', onnx_file)
+ continue
+ model = model_zoo.get_model(onnx_file)
+ if model.taskname not in self.models:
+ print('find model:', onnx_file, model.taskname)
+ self.models[model.taskname] = model
+ else:
+ print('duplicated model task type, ignore:', onnx_file, model.taskname)
+ del model
+ assert 'detection' in self.models
+ self.det_model = self.models['detection']
+
+
+ def prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640), mode ='None'):
+ self.det_thresh = det_thresh
+ self.mode = mode
+ assert det_size is not None
+ print('set det-size:', det_size)
+ self.det_size = det_size
+ for taskname, model in self.models.items():
+ if taskname=='detection':
+ model.prepare(ctx_id, input_size=det_size)
+ else:
+ model.prepare(ctx_id)
+
+ def get(self, img, crop_size, max_num=0):
+ bboxes, kpss = self.det_model.detect(img,
+ threshold=self.det_thresh,
+ max_num=max_num,
+ metric='default')
+ if bboxes.shape[0] == 0:
+ return None
+ ret = []
+ # for i in range(bboxes.shape[0]):
+ # bbox = bboxes[i, 0:4]
+ # det_score = bboxes[i, 4]
+ # kps = None
+ # if kpss is not None:
+ # kps = kpss[i]
+ # M, _ = face_align.estimate_norm(kps, crop_size, mode ='None')
+ # align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
+ align_img_list = []
+ M_list = []
+ for i in range(bboxes.shape[0]):
+ kps = None
+ if kpss is not None:
+ kps = kpss[i]
+ M, _ = face_align.estimate_norm(kps, crop_size, mode = self.mode)
+ align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
+ align_img_list.append(align_img)
+ M_list.append(M)
+
+ # det_score = bboxes[..., 4]
+
+ # best_index = np.argmax(det_score)
+
+ # kps = None
+ # if kpss is not None:
+ # kps = kpss[best_index]
+ # M, _ = face_align.estimate_norm(kps, crop_size, mode ='None')
+ # align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
+
+ return align_img_list, M_list
diff --git a/SimSwap/insightface_func/face_detect_crop_single.py b/SimSwap/insightface_func/face_detect_crop_single.py
new file mode 100644
index 0000000000000000000000000000000000000000..e60db36646ef04c640b95cc356431c3a5db79b82
--- /dev/null
+++ b/SimSwap/insightface_func/face_detect_crop_single.py
@@ -0,0 +1,91 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 17:03:58
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-24 16:46:04
+Description:
+'''
+from __future__ import division
+import collections
+import numpy as np
+import glob
+import os
+import os.path as osp
+import cv2
+from insightface.model_zoo import model_zoo
+from insightface_func.utils import face_align_ffhqandnewarc as face_align
+
+__all__ = ['Face_detect_crop', 'Face']
+
+Face = collections.namedtuple('Face', [
+ 'bbox', 'kps', 'det_score', 'embedding', 'gender', 'age',
+ 'embedding_norm', 'normed_embedding',
+ 'landmark'
+])
+
+Face.__new__.__defaults__ = (None, ) * len(Face._fields)
+
+
+class Face_detect_crop:
+ def __init__(self, name, root='~/.insightface_func/models'):
+ self.models = {}
+ root = os.path.expanduser(root)
+ onnx_files = glob.glob(osp.join(root, name, '*.onnx'))
+ onnx_files = sorted(onnx_files)
+ for onnx_file in onnx_files:
+ if onnx_file.find('_selfgen_')>0:
+ #print('ignore:', onnx_file)
+ continue
+ model = model_zoo.get_model(onnx_file)
+ if model.taskname not in self.models:
+ print('find model:', onnx_file, model.taskname)
+ self.models[model.taskname] = model
+ else:
+ print('duplicated model task type, ignore:', onnx_file, model.taskname)
+ del model
+ assert 'detection' in self.models
+ self.det_model = self.models['detection']
+
+
+ def prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640), mode ='None'):
+ self.det_thresh = det_thresh
+ self.mode = mode
+ assert det_size is not None
+ print('set det-size:', det_size)
+ self.det_size = det_size
+ for taskname, model in self.models.items():
+ if taskname=='detection':
+ model.prepare(ctx_id, input_size=det_size)
+ else:
+ model.prepare(ctx_id)
+
+
+ def get(self, img, crop_size, max_num=0):
+ bboxes, kpss = self.det_model.detect(
+ img,
+ threshold=self.det_thresh,
+ max_num=max_num,
+ metric='default'
+ )
+ if bboxes.shape[0] == 0:
+ return None
+
+ det_score = bboxes[..., 4]
+ best_index = np.argmax(det_score)
+
+ kps = None
+ if kpss is not None:
+ kps = kpss[best_index]
+
+ # 🔧 thêm margin dưới cằm
+ # dịch toàn bộ keypoints xuống 5% chiều cao khuôn mặt
+ h, w, _ = img.shape
+ face_h = bboxes[best_index][3] - bboxes[best_index][1]
+ kps[:,1] += int(0.1 * face_h) # tăng 0.1 = 10% chiều cao (giữ cằm)
+
+ # align về crop_size như cũ
+ M, _ = face_align.estimate_norm(kps, crop_size, mode=self.mode)
+ align_img = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0)
+
+ return [align_img], [M]
diff --git a/SimSwap/insightface_func/models/antelope/glintr100.onnx b/SimSwap/insightface_func/models/antelope/glintr100.onnx
new file mode 100644
index 0000000000000000000000000000000000000000..9d221846df998a9c85239fd74a9fe5685193775f
--- /dev/null
+++ b/SimSwap/insightface_func/models/antelope/glintr100.onnx
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ab1d6435d639628a6f3e5008dd4f929edf4c4124b1a7169e1048f9fef534cdf
+size 260665334
diff --git a/SimSwap/insightface_func/models/antelope/scrfd_10g_bnkps.onnx b/SimSwap/insightface_func/models/antelope/scrfd_10g_bnkps.onnx
new file mode 100644
index 0000000000000000000000000000000000000000..aa586e034379fa5ea5babc8aa73d47afcd0fa6c2
--- /dev/null
+++ b/SimSwap/insightface_func/models/antelope/scrfd_10g_bnkps.onnx
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5838f7fe053675b1c7a08b633df49e7af5495cee0493c7dcf6697200b85b5b91
+size 16923827
diff --git a/SimSwap/insightface_func/utils/__pycache__/face_align_ffhqandnewarc.cpython-310.pyc b/SimSwap/insightface_func/utils/__pycache__/face_align_ffhqandnewarc.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..215fa42fdb7cf89395b287b57914d7f53fbae81c
Binary files /dev/null and b/SimSwap/insightface_func/utils/__pycache__/face_align_ffhqandnewarc.cpython-310.pyc differ
diff --git a/SimSwap/insightface_func/utils/face_align_ffhqandnewarc.py b/SimSwap/insightface_func/utils/face_align_ffhqandnewarc.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f5aeb7977f78698955b909d94b9e0405aeaa5f5
--- /dev/null
+++ b/SimSwap/insightface_func/utils/face_align_ffhqandnewarc.py
@@ -0,0 +1,159 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-15 19:42:42
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-15 20:01:47
+Description:
+'''
+
+import cv2
+import numpy as np
+from skimage import transform as trans
+
+src1 = np.array([[51.642, 50.115], [57.617, 49.990], [35.740, 69.007],
+ [51.157, 89.050], [57.025, 89.702]],
+ dtype=np.float32)
+#<--left
+src2 = np.array([[45.031, 50.118], [65.568, 50.872], [39.677, 68.111],
+ [45.177, 86.190], [64.246, 86.758]],
+ dtype=np.float32)
+
+#---frontal
+src3 = np.array([[39.730, 51.138], [72.270, 51.138], [56.000, 68.493],
+ [42.463, 87.010], [69.537, 87.010]],
+ dtype=np.float32)
+
+#-->right
+src4 = np.array([[46.845, 50.872], [67.382, 50.118], [72.737, 68.111],
+ [48.167, 86.758], [67.236, 86.190]],
+ dtype=np.float32)
+
+#-->right profile
+src5 = np.array([[54.796, 49.990], [60.771, 50.115], [76.673, 69.007],
+ [55.388, 89.702], [61.257, 89.050]],
+ dtype=np.float32)
+
+src = np.array([src1, src2, src3, src4, src5])
+src_map = src
+
+ffhq_src = np.array([[192.98138, 239.94708], [318.90277, 240.1936], [256.63416, 314.01935],
+ [201.26117, 371.41043], [313.08905, 371.15118]])
+ffhq_src = np.expand_dims(ffhq_src, axis=0)
+
+# arcface_src = np.array(
+# [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366],
+# [41.5493, 92.3655], [70.7299, 92.2041]],
+# dtype=np.float32)
+
+# arcface_src = np.expand_dims(arcface_src, axis=0)
+
+# In[66]:
+
+
+# lmk is prediction; src is template
+def estimate_norm(lmk, image_size=112, mode='ffhq'):
+ assert lmk.shape == (5, 2)
+ tform = trans.SimilarityTransform()
+ lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1)
+ min_M = []
+ min_index = []
+ min_error = float('inf')
+ if mode == 'ffhq':
+ # assert image_size == 112
+ src = ffhq_src * image_size / 512
+ else:
+ src = src_map * image_size / 112
+ for i in np.arange(src.shape[0]):
+ tform.estimate(lmk, src[i])
+ M = tform.params[0:2, :]
+ results = np.dot(M, lmk_tran.T)
+ results = results.T
+ error = np.sum(np.sqrt(np.sum((results - src[i])**2, axis=1)))
+ # print(error)
+ if error < min_error:
+ min_error = error
+ min_M = M
+ min_index = i
+ return min_M, min_index
+
+
+def norm_crop(img, landmark, image_size=112, mode='ffhq'):
+ if mode == 'Both':
+ M_None, _ = estimate_norm(landmark, image_size, mode = 'newarc')
+ M_ffhq, _ = estimate_norm(landmark, image_size, mode='ffhq')
+ warped_None = cv2.warpAffine(img, M_None, (image_size, image_size), borderValue=0.0)
+ warped_ffhq = cv2.warpAffine(img, M_ffhq, (image_size, image_size), borderValue=0.0)
+ return warped_ffhq, warped_None
+ else:
+ M, pose_index = estimate_norm(landmark, image_size, mode)
+ warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)
+ return warped
+
+def square_crop(im, S):
+ if im.shape[0] > im.shape[1]:
+ height = S
+ width = int(float(im.shape[1]) / im.shape[0] * S)
+ scale = float(S) / im.shape[0]
+ else:
+ width = S
+ height = int(float(im.shape[0]) / im.shape[1] * S)
+ scale = float(S) / im.shape[1]
+ resized_im = cv2.resize(im, (width, height))
+ det_im = np.zeros((S, S, 3), dtype=np.uint8)
+ det_im[:resized_im.shape[0], :resized_im.shape[1], :] = resized_im
+ return det_im, scale
+
+
+def transform(data, center, output_size, scale, rotation):
+ scale_ratio = scale
+ rot = float(rotation) * np.pi / 180.0
+ #translation = (output_size/2-center[0]*scale_ratio, output_size/2-center[1]*scale_ratio)
+ t1 = trans.SimilarityTransform(scale=scale_ratio)
+ cx = center[0] * scale_ratio
+ cy = center[1] * scale_ratio
+ t2 = trans.SimilarityTransform(translation=(-1 * cx, -1 * cy))
+ t3 = trans.SimilarityTransform(rotation=rot)
+ t4 = trans.SimilarityTransform(translation=(output_size / 2,
+ output_size / 2))
+ t = t1 + t2 + t3 + t4
+ M = t.params[0:2]
+ cropped = cv2.warpAffine(data,
+ M, (output_size, output_size),
+ borderValue=0.0)
+ return cropped, M
+
+
+def trans_points2d(pts, M):
+ new_pts = np.zeros(shape=pts.shape, dtype=np.float32)
+ for i in range(pts.shape[0]):
+ pt = pts[i]
+ new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32)
+ new_pt = np.dot(M, new_pt)
+ #print('new_pt', new_pt.shape, new_pt)
+ new_pts[i] = new_pt[0:2]
+
+ return new_pts
+
+
+def trans_points3d(pts, M):
+ scale = np.sqrt(M[0][0] * M[0][0] + M[0][1] * M[0][1])
+ #print(scale)
+ new_pts = np.zeros(shape=pts.shape, dtype=np.float32)
+ for i in range(pts.shape[0]):
+ pt = pts[i]
+ new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32)
+ new_pt = np.dot(M, new_pt)
+ #print('new_pt', new_pt.shape, new_pt)
+ new_pts[i][0:2] = new_pt[0:2]
+ new_pts[i][2] = pts[i][2] * scale
+
+ return new_pts
+
+
+def trans_points(pts, M):
+ if pts.shape[1] == 2:
+ return trans_points2d(pts, M)
+ else:
+ return trans_points3d(pts, M)
+
diff --git a/SimSwap/models/__init__.py b/SimSwap/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..289de91be1d50609021a09d84130981cf4da9bb8
--- /dev/null
+++ b/SimSwap/models/__init__.py
@@ -0,0 +1,4 @@
+from .arcface_models import ArcMarginModel
+from .arcface_models import ResNet
+from .arcface_models import IRBlock
+from .arcface_models import SEBlock
\ No newline at end of file
diff --git a/SimSwap/models/__pycache__/__init__.cpython-310.pyc b/SimSwap/models/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..86701a5a42a8f90ef4ee022ad990676157c8c90e
Binary files /dev/null and b/SimSwap/models/__pycache__/__init__.cpython-310.pyc differ
diff --git a/SimSwap/models/__pycache__/arcface_models.cpython-310.pyc b/SimSwap/models/__pycache__/arcface_models.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f81003dd82a52fccfe73a17584b71f85fa62197c
Binary files /dev/null and b/SimSwap/models/__pycache__/arcface_models.cpython-310.pyc differ
diff --git a/SimSwap/models/__pycache__/base_model.cpython-310.pyc b/SimSwap/models/__pycache__/base_model.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8d3dbd61a43f666efeead465e6fa88e849e523c5
Binary files /dev/null and b/SimSwap/models/__pycache__/base_model.cpython-310.pyc differ
diff --git a/SimSwap/models/__pycache__/config.cpython-310.pyc b/SimSwap/models/__pycache__/config.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..47f200988765e43afc0d96a317ef4254b78c2a92
Binary files /dev/null and b/SimSwap/models/__pycache__/config.cpython-310.pyc differ
diff --git a/SimSwap/models/__pycache__/fs_model.cpython-310.pyc b/SimSwap/models/__pycache__/fs_model.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..07289d740a15f2e0b4638057f66d7dac0440bec8
Binary files /dev/null and b/SimSwap/models/__pycache__/fs_model.cpython-310.pyc differ
diff --git a/SimSwap/models/__pycache__/fs_networks.cpython-310.pyc b/SimSwap/models/__pycache__/fs_networks.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4c18ddefc204f83002a0e6fd0e42ac274e0439f2
Binary files /dev/null and b/SimSwap/models/__pycache__/fs_networks.cpython-310.pyc differ
diff --git a/SimSwap/models/__pycache__/fs_networks_512.cpython-310.pyc b/SimSwap/models/__pycache__/fs_networks_512.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0b54f1181d800e45bf61509ee731e1dc3ca49ab7
Binary files /dev/null and b/SimSwap/models/__pycache__/fs_networks_512.cpython-310.pyc differ
diff --git a/SimSwap/models/__pycache__/models.cpython-310.pyc b/SimSwap/models/__pycache__/models.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8f503bbbebf823188d3c29bc6702e5988879bbda
Binary files /dev/null and b/SimSwap/models/__pycache__/models.cpython-310.pyc differ
diff --git a/SimSwap/models/__pycache__/networks.cpython-310.pyc b/SimSwap/models/__pycache__/networks.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..79a7f15327170bafca7703304520e074a9587bd3
Binary files /dev/null and b/SimSwap/models/__pycache__/networks.cpython-310.pyc differ
diff --git a/SimSwap/models/arcface_models.py b/SimSwap/models/arcface_models.py
new file mode 100644
index 0000000000000000000000000000000000000000..39a6ac54e605defacef85932c036fe317272e316
--- /dev/null
+++ b/SimSwap/models/arcface_models.py
@@ -0,0 +1,163 @@
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import Parameter
+from .config import device, num_classes
+
+
+
+class SEBlock(nn.Module):
+ def __init__(self, channel, reduction=16):
+ super(SEBlock, self).__init__()
+ self.avg_pool = nn.AdaptiveAvgPool2d(1)
+ self.fc = nn.Sequential(
+ nn.Linear(channel, channel // reduction),
+ nn.PReLU(),
+ nn.Linear(channel // reduction, channel),
+ nn.Sigmoid()
+ )
+
+ def forward(self, x):
+ b, c, _, _ = x.size()
+ y = self.avg_pool(x).view(b, c)
+ y = self.fc(y).view(b, c, 1, 1)
+ return x * y
+
+
+class IRBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True):
+ super(IRBlock, self).__init__()
+ self.bn0 = nn.BatchNorm2d(inplanes)
+ self.conv1 = conv3x3(inplanes, inplanes)
+ self.bn1 = nn.BatchNorm2d(inplanes)
+ self.prelu = nn.PReLU()
+ self.conv2 = conv3x3(inplanes, planes, stride)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.downsample = downsample
+ self.stride = stride
+ self.use_se = use_se
+ if self.use_se:
+ self.se = SEBlock(planes)
+
+ def forward(self, x):
+ residual = x
+ out = self.bn0(x)
+ out = self.conv1(out)
+ out = self.bn1(out)
+ out = self.prelu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ if self.use_se:
+ out = self.se(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.prelu(out)
+
+ return out
+
+
+class ResNet(nn.Module):
+
+ def __init__(self, block, layers, use_se=True):
+ self.inplanes = 64
+ self.use_se = use_se
+ super(ResNet, self).__init__()
+ self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(64)
+ self.prelu = nn.PReLU()
+ self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
+ self.layer1 = self._make_layer(block, 64, layers[0])
+ self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+ self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+ self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+ self.bn2 = nn.BatchNorm2d(512)
+ self.dropout = nn.Dropout()
+ self.fc = nn.Linear(512 * 7 * 7, 512)
+ self.bn3 = nn.BatchNorm1d(512)
+
+ for m in self.modules():
+ if isinstance(m, nn.Conv2d):
+ nn.init.xavier_normal_(m.weight)
+ elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+ elif isinstance(m, nn.Linear):
+ nn.init.xavier_normal_(m.weight)
+ nn.init.constant_(m.bias, 0)
+
+ def _make_layer(self, block, planes, blocks, stride=1):
+ downsample = None
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ nn.Conv2d(self.inplanes, planes * block.expansion,
+ kernel_size=1, stride=stride, bias=False),
+ nn.BatchNorm2d(planes * block.expansion),
+ )
+
+ layers = []
+ layers.append(block(self.inplanes, planes, stride, downsample, use_se=self.use_se))
+ self.inplanes = planes
+ for i in range(1, blocks):
+ layers.append(block(self.inplanes, planes, use_se=self.use_se))
+
+ return nn.Sequential(*layers)
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.prelu(x)
+ x = self.maxpool(x)
+
+ x = self.layer1(x)
+ x = self.layer2(x)
+ x = self.layer3(x)
+ x = self.layer4(x)
+
+ x = self.bn2(x)
+ x = self.dropout(x)
+ # feature = x
+ x = x.view(x.size(0), -1)
+ x = self.fc(x)
+ x = self.bn3(x)
+
+ return x
+
+
+class ArcMarginModel(nn.Module):
+ def __init__(self, args):
+ super(ArcMarginModel, self).__init__()
+
+ self.weight = Parameter(torch.FloatTensor(num_classes, args.emb_size))
+ nn.init.xavier_uniform_(self.weight)
+
+ self.easy_margin = args.easy_margin
+ self.m = args.margin_m
+ self.s = args.margin_s
+
+ self.cos_m = math.cos(self.m)
+ self.sin_m = math.sin(self.m)
+ self.th = math.cos(math.pi - self.m)
+ self.mm = math.sin(math.pi - self.m) * self.m
+
+ def forward(self, input, label):
+ x = F.normalize(input)
+ W = F.normalize(self.weight)
+ cosine = F.linear(x, W)
+ sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
+ phi = cosine * self.cos_m - sine * self.sin_m # cos(theta + m)
+ if self.easy_margin:
+ phi = torch.where(cosine > 0, phi, cosine)
+ else:
+ phi = torch.where(cosine > self.th, phi, cosine - self.mm)
+ one_hot = torch.zeros(cosine.size(), device=device)
+ one_hot.scatter_(1, label.view(-1, 1).long(), 1)
+ output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
+ output *= self.s
+ return output
\ No newline at end of file
diff --git a/SimSwap/models/base_model.py b/SimSwap/models/base_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..0a6474afe364de9c269b3080cbccab0036797230
--- /dev/null
+++ b/SimSwap/models/base_model.py
@@ -0,0 +1,138 @@
+import os
+import torch
+import sys
+
+class BaseModel(torch.nn.Module):
+ def name(self):
+ return 'BaseModel'
+
+ def initialize(self, opt):
+ self.opt = opt
+ self.gpu_ids = opt.gpu_ids
+ self.isTrain = opt.isTrain
+ self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor
+ self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+
+ def set_input(self, input):
+ self.input = input
+
+ def forward(self):
+ pass
+
+ # used in test time, no backprop
+ def test(self):
+ pass
+
+ def get_image_paths(self):
+ pass
+
+ def optimize_parameters(self):
+ pass
+
+ def get_current_visuals(self):
+ return self.input
+
+ def get_current_errors(self):
+ return {}
+
+ def save(self, label):
+ pass
+
+ # helper saving function that can be used by subclasses
+ def save_network(self, network, network_label, epoch_label, gpu_ids=None):
+ save_filename = '{}_net_{}.pth'.format(epoch_label, network_label)
+ save_path = os.path.join(self.save_dir, save_filename)
+ torch.save(network.cpu().state_dict(), save_path)
+ if torch.cuda.is_available():
+ network.cuda()
+
+ def save_optim(self, network, network_label, epoch_label, gpu_ids=None):
+ save_filename = '{}_optim_{}.pth'.format(epoch_label, network_label)
+ save_path = os.path.join(self.save_dir, save_filename)
+ torch.save(network.state_dict(), save_path)
+
+
+ # helper loading function that can be used by subclasses
+ def load_network(self, network, network_label, epoch_label, save_dir=''):
+ save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
+ if not save_dir:
+ save_dir = self.save_dir
+ save_path = os.path.join(save_dir, save_filename)
+ if not os.path.isfile(save_path):
+ print('%s not exists yet!' % save_path)
+ if network_label == 'G':
+ raise('Generator must exist!')
+ else:
+ #network.load_state_dict(torch.load(save_path))
+ try:
+ network.load_state_dict(torch.load(save_path))
+ except:
+ pretrained_dict = torch.load(save_path)
+ model_dict = network.state_dict()
+ try:
+ pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+ network.load_state_dict(pretrained_dict)
+ if self.opt.verbose:
+ print('Pretrained network %s has excessive layers; Only loading layers that are used' % network_label)
+ except:
+ print('Pretrained network %s has fewer layers; The following are not initialized:' % network_label)
+ for k, v in pretrained_dict.items():
+ if v.size() == model_dict[k].size():
+ model_dict[k] = v
+
+ if sys.version_info >= (3,0):
+ not_initialized = set()
+ else:
+ from sets import Set
+ not_initialized = Set()
+
+ for k, v in model_dict.items():
+ if k not in pretrained_dict or v.size() != pretrained_dict[k].size():
+ not_initialized.add(k.split('.')[0])
+
+ print(sorted(not_initialized))
+ network.load_state_dict(model_dict)
+
+ # helper loading function that can be used by subclasses
+ def load_optim(self, network, network_label, epoch_label, save_dir=''):
+ save_filename = '%s_optim_%s.pth' % (epoch_label, network_label)
+ if not save_dir:
+ save_dir = self.save_dir
+ save_path = os.path.join(save_dir, save_filename)
+ if not os.path.isfile(save_path):
+ print('%s not exists yet!' % save_path)
+ if network_label == 'G':
+ raise('Generator must exist!')
+ else:
+ #network.load_state_dict(torch.load(save_path))
+ try:
+ network.load_state_dict(torch.load(save_path, map_location=torch.device("cpu")))
+ except:
+ pretrained_dict = torch.load(save_path, map_location=torch.device("cpu"))
+ model_dict = network.state_dict()
+ try:
+ pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+ network.load_state_dict(pretrained_dict)
+ if self.opt.verbose:
+ print('Pretrained network %s has excessive layers; Only loading layers that are used' % network_label)
+ except:
+ print('Pretrained network %s has fewer layers; The following are not initialized:' % network_label)
+ for k, v in pretrained_dict.items():
+ if v.size() == model_dict[k].size():
+ model_dict[k] = v
+
+ if sys.version_info >= (3,0):
+ not_initialized = set()
+ else:
+ from sets import Set
+ not_initialized = Set()
+
+ for k, v in model_dict.items():
+ if k not in pretrained_dict or v.size() != pretrained_dict[k].size():
+ not_initialized.add(k.split('.')[0])
+
+ print(sorted(not_initialized))
+ network.load_state_dict(model_dict)
+
+ def update_learning_rate():
+ pass
diff --git a/SimSwap/models/config.py b/SimSwap/models/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb83edb1152974ef9d7ff7c46d3b98ffbe9b2495
--- /dev/null
+++ b/SimSwap/models/config.py
@@ -0,0 +1,28 @@
+import os
+
+import torch
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # sets device for model and PyTorch tensors
+
+# Model parameters
+image_w = 112
+image_h = 112
+channel = 3
+emb_size = 512
+
+# Training parameters
+num_workers = 1 # for data-loading; right now, only 1 works with h5py
+grad_clip = 5. # clip gradients at an absolute value of
+print_freq = 100 # print training/validation stats every __ batches
+checkpoint = None # path to checkpoint, None if none
+
+# Data parameters
+num_classes = 93431
+num_samples = 5179510
+DATA_DIR = 'data'
+# faces_ms1m_folder = 'data/faces_ms1m_112x112'
+faces_ms1m_folder = 'data/ms1m-retinaface-t1'
+path_imgidx = os.path.join(faces_ms1m_folder, 'train.idx')
+path_imgrec = os.path.join(faces_ms1m_folder, 'train.rec')
+IMG_DIR = 'data/images'
+pickle_file = 'data/faces_ms1m_112x112.pickle'
diff --git a/SimSwap/models/fs_model.py b/SimSwap/models/fs_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..95ac877147d580605be5c5ead600f0edb6371520
--- /dev/null
+++ b/SimSwap/models/fs_model.py
@@ -0,0 +1,242 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import os
+from torch.autograd import Variable
+from .base_model import BaseModel
+from . import networks
+
+class SpecificNorm(nn.Module):
+ def __init__(self, epsilon=1e-8):
+ """
+ @notice: avoid in-place ops.
+ https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
+ """
+ super(SpecificNorm, self).__init__()
+ self.mean = np.array([0.485, 0.456, 0.406])
+ self.mean = torch.from_numpy(self.mean).float().cuda()
+ self.mean = self.mean.view([1, 3, 1, 1])
+
+ self.std = np.array([0.229, 0.224, 0.225])
+ self.std = torch.from_numpy(self.std).float().cuda()
+ self.std = self.std.view([1, 3, 1, 1])
+
+ def forward(self, x):
+ mean = self.mean.expand([1, 3, x.shape[2], x.shape[3]])
+ std = self.std.expand([1, 3, x.shape[2], x.shape[3]])
+
+ x = (x - mean) / std
+
+ return x
+
+class fsModel(BaseModel):
+ def name(self):
+ return 'fsModel'
+
+ def init_loss_filter(self, use_gan_feat_loss, use_vgg_loss):
+ flags = (True, use_gan_feat_loss, use_vgg_loss, True, True, True, True, True)
+
+ def loss_filter(g_gan, g_gan_feat, g_vgg, g_id, g_rec, g_mask, d_real, d_fake):
+ return [l for (l, f) in zip((g_gan, g_gan_feat, g_vgg, g_id, g_rec, g_mask, d_real, d_fake), flags) if f]
+
+ return loss_filter
+
+ def initialize(self, opt):
+ BaseModel.initialize(self, opt)
+ if opt.resize_or_crop != 'none' or not opt.isTrain: # when training at full res this causes OOM
+ torch.backends.cudnn.benchmark = True
+ self.isTrain = opt.isTrain
+
+ device = torch.device("cuda:0")
+
+ if opt.crop_size == 224:
+ from .fs_networks import Generator_Adain_Upsample, Discriminator
+ elif opt.crop_size == 512:
+ from .fs_networks_512 import Generator_Adain_Upsample, Discriminator
+
+ # Generator network
+ self.netG = Generator_Adain_Upsample(input_nc=3, output_nc=3, latent_size=512, n_blocks=9, deep=False)
+ self.netG.to(device)
+
+ # Id network
+ netArc_checkpoint = opt.Arc_path
+ netArc_checkpoint = torch.load(netArc_checkpoint, map_location=torch.device("cpu"))
+ self.netArc = netArc_checkpoint
+ self.netArc = self.netArc.to(device)
+ self.netArc.eval()
+
+ if not self.isTrain:
+ pretrained_path = '' if not self.isTrain else opt.load_pretrain
+ self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)
+ return
+
+ # Discriminator network
+ if opt.gan_mode == 'original':
+ use_sigmoid = True
+ else:
+ use_sigmoid = False
+ self.netD1 = Discriminator(input_nc=3, use_sigmoid=use_sigmoid)
+ self.netD2 = Discriminator(input_nc=3, use_sigmoid=use_sigmoid)
+ self.netD1.to(device)
+ self.netD2.to(device)
+
+ #
+ self.spNorm =SpecificNorm()
+ self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
+
+ # load networks
+ if opt.continue_train or opt.load_pretrain:
+ pretrained_path = '' if not self.isTrain else opt.load_pretrain
+ # print (pretrained_path)
+ self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)
+ self.load_network(self.netD1, 'D1', opt.which_epoch, pretrained_path)
+ self.load_network(self.netD2, 'D2', opt.which_epoch, pretrained_path)
+
+
+
+ if self.isTrain:
+ # define loss functions
+ self.loss_filter = self.init_loss_filter(not opt.no_ganFeat_loss, not opt.no_vgg_loss)
+
+ self.criterionGAN = networks.GANLoss(opt.gan_mode, tensor=self.Tensor, opt=self.opt)
+ self.criterionFeat = nn.L1Loss()
+ self.criterionRec = nn.L1Loss()
+
+ # Names so we can breakout loss
+ self.loss_names = self.loss_filter('G_GAN', 'G_GAN_Feat', 'G_VGG', 'G_ID', 'G_Rec', 'D_GP',
+ 'D_real', 'D_fake')
+
+ # initialize optimizers
+
+ # optimizer G
+ params = list(self.netG.parameters())
+ self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+
+ # optimizer D
+ params = list(self.netD1.parameters()) + list(self.netD2.parameters())
+ self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+
+ def _gradinet_penalty_D(self, netD, img_att, img_fake):
+ # interpolate sample
+ bs = img_fake.shape[0]
+ alpha = torch.rand(bs, 1, 1, 1).expand_as(img_fake).cuda()
+ interpolated = Variable(alpha * img_att + (1 - alpha) * img_fake, requires_grad=True)
+ pred_interpolated = netD.forward(interpolated)
+ pred_interpolated = pred_interpolated[-1]
+
+ # compute gradients
+ grad = torch.autograd.grad(outputs=pred_interpolated,
+ inputs=interpolated,
+ grad_outputs=torch.ones(pred_interpolated.size()).cuda(),
+ retain_graph=True,
+ create_graph=True,
+ only_inputs=True)[0]
+
+ # penalize gradients
+ grad = grad.view(grad.size(0), -1)
+ grad_l2norm = torch.sqrt(torch.sum(grad ** 2, dim=1))
+ loss_d_gp = torch.mean((grad_l2norm - 1) ** 2)
+
+ return loss_d_gp
+
+ def cosin_metric(self, x1, x2):
+ #return np.dot(x1, x2) / (np.linalg.norm(x1) * np.linalg.norm(x2))
+ return torch.sum(x1 * x2, dim=1) / (torch.norm(x1, dim=1) * torch.norm(x2, dim=1))
+
+ def forward(self, img_id, img_att, latent_id, latent_att, for_G=False):
+ loss_D_fake, loss_D_real, loss_D_GP = 0, 0, 0
+ loss_G_GAN, loss_G_GAN_Feat, loss_G_VGG, loss_G_ID, loss_G_Rec = 0,0,0,0,0
+
+ img_fake = self.netG.forward(img_att, latent_id)
+ if not self.isTrain:
+ return img_fake
+ img_fake_downsample = self.downsample(img_fake)
+ img_att_downsample = self.downsample(img_att)
+
+
+
+ # D_Fake
+ fea1_fake = self.netD1.forward(img_fake.detach())
+ fea2_fake = self.netD2.forward(img_fake_downsample.detach())
+ pred_fake = [fea1_fake, fea2_fake]
+ loss_D_fake = self.criterionGAN(pred_fake, False, for_discriminator=True)
+
+
+ # D_Feal
+ fea1_real = self.netD1.forward(img_att)
+ fea2_real = self.netD2.forward(img_att_downsample)
+ pred_real = [fea1_real, fea2_real]
+ fea_real = [fea1_real, fea2_real]
+ loss_D_real = self.criterionGAN(pred_real, True, for_discriminator=True)
+ #print('=====================D_Real========================')
+
+ # D_GP
+
+ loss_D_GP = 0
+
+ # G_GAN
+ fea1_fake = self.netD1.forward(img_fake)
+ fea2_fake = self.netD2.forward(img_fake_downsample)
+ #pred_fake = [fea1_fake[-1], fea2_fake[-1]]
+ pred_fake = [fea1_fake, fea2_fake]
+ fea_fake = [fea1_fake, fea2_fake]
+ loss_G_GAN = self.criterionGAN(pred_fake, True, for_discriminator=False)
+
+ # GAN feature matching loss
+ n_layers_D = 4
+ num_D = 2
+ if not self.opt.no_ganFeat_loss:
+ feat_weights = 4.0 / (n_layers_D + 1)
+ D_weights = 1.0 / num_D
+ for i in range(num_D):
+ for j in range(0, len(fea_fake[i]) - 1):
+ loss_G_GAN_Feat += D_weights * feat_weights * \
+ self.criterionFeat(fea_fake[i][j],
+ fea_real[i][j].detach()) * self.opt.lambda_feat
+
+
+ #G_ID
+ img_fake_down = F.interpolate(img_fake, size=(112,112))
+ img_fake_down = self.spNorm(img_fake_down)
+ latent_fake = self.netArc(img_fake_down)
+ loss_G_ID = (1 - self.cosin_metric(latent_fake, latent_id))
+ #print('=====================G_ID========================')
+ #print(loss_G_ID)
+
+ #G_Rec
+ loss_G_Rec = self.criterionRec(img_fake, img_att) * self.opt.lambda_rec
+
+ # Only return the fake_B image if necessary to save BW
+ return [self.loss_filter(loss_G_GAN, loss_G_GAN_Feat, loss_G_VGG, loss_G_ID, loss_G_Rec, loss_D_GP, loss_D_real, loss_D_fake),
+ img_fake]
+
+
+ def save(self, which_epoch):
+ self.save_network(self.netG, 'G', which_epoch, self.gpu_ids)
+ self.save_network(self.netD1, 'D1', which_epoch, self.gpu_ids)
+ self.save_network(self.netD2, 'D2', which_epoch, self.gpu_ids)
+ '''if self.gen_features:
+ self.save_network(self.netE, 'E', which_epoch, self.gpu_ids)'''
+
+ def update_fixed_params(self):
+ # after fixing the global generator for a number of iterations, also start finetuning it
+ params = list(self.netG.parameters())
+ if self.gen_features:
+ params += list(self.netE.parameters())
+ self.optimizer_G = torch.optim.Adam(params, lr=self.opt.lr, betas=(self.opt.beta1, 0.999))
+ if self.opt.verbose:
+ print('------------ Now also finetuning global generator -----------')
+
+ def update_learning_rate(self):
+ lrd = self.opt.lr / self.opt.niter_decay
+ lr = self.old_lr - lrd
+ for param_group in self.optimizer_D.param_groups:
+ param_group['lr'] = lr
+ for param_group in self.optimizer_G.param_groups:
+ param_group['lr'] = lr
+ if self.opt.verbose:
+ print('update learning rate: %f -> %f' % (self.old_lr, lr))
+ self.old_lr = lr
+
+
diff --git a/SimSwap/models/fs_networks.py b/SimSwap/models/fs_networks.py
new file mode 100644
index 0000000000000000000000000000000000000000..227f6faba9c786f502544e323b9ea77e2b47613c
--- /dev/null
+++ b/SimSwap/models/fs_networks.py
@@ -0,0 +1,215 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation. All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import torch
+import torch.nn as nn
+
+
+class InstanceNorm(nn.Module):
+ def __init__(self, epsilon=1e-8):
+ """
+ @notice: avoid in-place ops.
+ https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
+ """
+ super(InstanceNorm, self).__init__()
+ self.epsilon = epsilon
+
+ def forward(self, x):
+ x = x - torch.mean(x, (2, 3), True)
+ tmp = torch.mul(x, x) # or x ** 2
+ tmp = torch.rsqrt(torch.mean(tmp, (2, 3), True) + self.epsilon)
+ return x * tmp
+
+class ApplyStyle(nn.Module):
+ """
+ @ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb
+ """
+ def __init__(self, latent_size, channels):
+ super(ApplyStyle, self).__init__()
+ self.linear = nn.Linear(latent_size, channels * 2)
+
+ def forward(self, x, latent):
+ style = self.linear(latent) # style => [batch_size, n_channels*2]
+ shape = [-1, 2, x.size(1), 1, 1]
+ style = style.view(shape) # [batch_size, 2, n_channels, ...]
+ #x = x * (style[:, 0] + 1.) + style[:, 1]
+ x = x * (style[:, 0] * 1 + 1.) + style[:, 1] * 1
+ return x
+
+class ResnetBlock_Adain(nn.Module):
+ def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True)):
+ super(ResnetBlock_Adain, self).__init__()
+
+ p = 0
+ conv1 = []
+ if padding_type == 'reflect':
+ conv1 += [nn.ReflectionPad2d(1)]
+ elif padding_type == 'replicate':
+ conv1 += [nn.ReplicationPad2d(1)]
+ elif padding_type == 'zero':
+ p = 1
+ else:
+ raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+ conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), InstanceNorm()]
+ self.conv1 = nn.Sequential(*conv1)
+ self.style1 = ApplyStyle(latent_size, dim)
+ self.act1 = activation
+
+ p = 0
+ conv2 = []
+ if padding_type == 'reflect':
+ conv2 += [nn.ReflectionPad2d(1)]
+ elif padding_type == 'replicate':
+ conv2 += [nn.ReplicationPad2d(1)]
+ elif padding_type == 'zero':
+ p = 1
+ else:
+ raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+ conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), InstanceNorm()]
+ self.conv2 = nn.Sequential(*conv2)
+ self.style2 = ApplyStyle(latent_size, dim)
+
+
+ def forward(self, x, dlatents_in_slice):
+ y = self.conv1(x)
+ y = self.style1(y, dlatents_in_slice)
+ y = self.act1(y)
+ y = self.conv2(y)
+ y = self.style2(y, dlatents_in_slice)
+ out = x + y
+ return out
+
+
+
+class Generator_Adain_Upsample(nn.Module):
+ def __init__(self, input_nc, output_nc, latent_size, n_blocks=6, deep=False,
+ norm_layer=nn.BatchNorm2d,
+ padding_type='reflect'):
+ assert (n_blocks >= 0)
+ super(Generator_Adain_Upsample, self).__init__()
+ activation = nn.ReLU(True)
+ self.deep = deep
+
+ self.first_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(input_nc, 64, kernel_size=7, padding=0),
+ norm_layer(64), activation)
+ ### downsample
+ self.down1 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
+ norm_layer(128), activation)
+ self.down2 = nn.Sequential(nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
+ norm_layer(256), activation)
+ self.down3 = nn.Sequential(nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),
+ norm_layer(512), activation)
+ if self.deep:
+ self.down4 = nn.Sequential(nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),
+ norm_layer(512), activation)
+
+ ### resnet blocks
+ BN = []
+ for i in range(n_blocks):
+ BN += [
+ ResnetBlock_Adain(512, latent_size=latent_size, padding_type=padding_type, activation=activation)]
+ self.BottleNeck = nn.Sequential(*BN)
+
+ if self.deep:
+ self.up4 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear'),
+ nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(512), activation
+ )
+ self.up3 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear'),
+ nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(256), activation
+ )
+ self.up2 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear'),
+ nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(128), activation
+ )
+ self.up1 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear'),
+ nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(64), activation
+ )
+ self.last_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(64, output_nc, kernel_size=7, padding=0),
+ nn.Tanh())
+
+ def forward(self, input, dlatents):
+ x = input # 3*224*224
+
+ skip1 = self.first_layer(x)
+ skip2 = self.down1(skip1)
+ skip3 = self.down2(skip2)
+ if self.deep:
+ skip4 = self.down3(skip3)
+ x = self.down4(skip4)
+ else:
+ x = self.down3(skip3)
+
+ for i in range(len(self.BottleNeck)):
+ x = self.BottleNeck[i](x, dlatents)
+
+ if self.deep:
+ x = self.up4(x)
+ x = self.up3(x)
+ x = self.up2(x)
+ x = self.up1(x)
+ x = self.last_layer(x)
+ x = (x + 1) / 2
+
+ return x
+
+class Discriminator(nn.Module):
+ def __init__(self, input_nc, norm_layer=nn.BatchNorm2d, use_sigmoid=False):
+ super(Discriminator, self).__init__()
+
+ kw = 4
+ padw = 1
+ self.down1 = nn.Sequential(
+ nn.Conv2d(input_nc, 64, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)
+ )
+ self.down2 = nn.Sequential(
+ nn.Conv2d(64, 128, kernel_size=kw, stride=2, padding=padw),
+ norm_layer(128), nn.LeakyReLU(0.2, True)
+ )
+ self.down3 = nn.Sequential(
+ nn.Conv2d(128, 256, kernel_size=kw, stride=2, padding=padw),
+ norm_layer(256), nn.LeakyReLU(0.2, True)
+ )
+ self.down4 = nn.Sequential(
+ nn.Conv2d(256, 512, kernel_size=kw, stride=2, padding=padw),
+ norm_layer(512), nn.LeakyReLU(0.2, True)
+ )
+ self.conv1 = nn.Sequential(
+ nn.Conv2d(512, 512, kernel_size=kw, stride=1, padding=padw),
+ norm_layer(512),
+ nn.LeakyReLU(0.2, True)
+ )
+
+ if use_sigmoid:
+ self.conv2 = nn.Sequential(
+ nn.Conv2d(512, 1, kernel_size=kw, stride=1, padding=padw), nn.Sigmoid()
+ )
+ else:
+ self.conv2 = nn.Sequential(
+ nn.Conv2d(512, 1, kernel_size=kw, stride=1, padding=padw)
+ )
+
+ def forward(self, input):
+ out = []
+ x = self.down1(input)
+ out.append(x)
+ x = self.down2(x)
+ out.append(x)
+ x = self.down3(x)
+ out.append(x)
+ x = self.down4(x)
+ out.append(x)
+ x = self.conv1(x)
+ out.append(x)
+ x = self.conv2(x)
+ out.append(x)
+
+ return out
\ No newline at end of file
diff --git a/SimSwap/models/fs_networks_512.py b/SimSwap/models/fs_networks_512.py
new file mode 100644
index 0000000000000000000000000000000000000000..1b9f9c8f39cd2eb34c4d63b90e3d1b1052d53482
--- /dev/null
+++ b/SimSwap/models/fs_networks_512.py
@@ -0,0 +1,232 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 16:55:48
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-24 16:58:06
+Description:
+'''
+"""
+Copyright (C) 2019 NVIDIA Corporation. All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import torch
+import torch.nn as nn
+
+
+class InstanceNorm(nn.Module):
+ def __init__(self, epsilon=1e-8):
+ """
+ @notice: avoid in-place ops.
+ https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
+ """
+ super(InstanceNorm, self).__init__()
+ self.epsilon = epsilon
+
+ def forward(self, x):
+ x = x - torch.mean(x, (2, 3), True)
+ tmp = torch.mul(x, x) # or x ** 2
+ tmp = torch.rsqrt(torch.mean(tmp, (2, 3), True) + self.epsilon)
+ return x * tmp
+
+class ApplyStyle(nn.Module):
+ """
+ @ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb
+ """
+ def __init__(self, latent_size, channels):
+ super(ApplyStyle, self).__init__()
+ self.linear = nn.Linear(latent_size, channels * 2)
+
+ def forward(self, x, latent):
+ style = self.linear(latent) # style => [batch_size, n_channels*2]
+ shape = [-1, 2, x.size(1), 1, 1]
+ style = style.view(shape) # [batch_size, 2, n_channels, ...]
+ #x = x * (style[:, 0] + 1.) + style[:, 1]
+ x = x * (style[:, 0] * 1 + 1.) + style[:, 1] * 1
+ return x
+
+class ResnetBlock_Adain(nn.Module):
+ def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True)):
+ super(ResnetBlock_Adain, self).__init__()
+
+ p = 0
+ conv1 = []
+ if padding_type == 'reflect':
+ conv1 += [nn.ReflectionPad2d(1)]
+ elif padding_type == 'replicate':
+ conv1 += [nn.ReplicationPad2d(1)]
+ elif padding_type == 'zero':
+ p = 1
+ else:
+ raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+ conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), InstanceNorm()]
+ self.conv1 = nn.Sequential(*conv1)
+ self.style1 = ApplyStyle(latent_size, dim)
+ self.act1 = activation
+
+ p = 0
+ conv2 = []
+ if padding_type == 'reflect':
+ conv2 += [nn.ReflectionPad2d(1)]
+ elif padding_type == 'replicate':
+ conv2 += [nn.ReplicationPad2d(1)]
+ elif padding_type == 'zero':
+ p = 1
+ else:
+ raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+ conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), InstanceNorm()]
+ self.conv2 = nn.Sequential(*conv2)
+ self.style2 = ApplyStyle(latent_size, dim)
+
+
+ def forward(self, x, dlatents_in_slice):
+ y = self.conv1(x)
+ y = self.style1(y, dlatents_in_slice)
+ y = self.act1(y)
+ y = self.conv2(y)
+ y = self.style2(y, dlatents_in_slice)
+ out = x + y
+ return out
+
+
+
+class Generator_Adain_Upsample(nn.Module):
+ def __init__(self, input_nc, output_nc, latent_size, n_blocks=6, deep=False,
+ norm_layer=nn.BatchNorm2d,
+ padding_type='reflect'):
+ assert (n_blocks >= 0)
+ super(Generator_Adain_Upsample, self).__init__()
+ activation = nn.ReLU(True)
+ self.deep = deep
+
+ self.first_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(input_nc, 32, kernel_size=7, padding=0),
+ norm_layer(32), activation)
+ ### downsample
+ self.down0 = nn.Sequential(nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
+ norm_layer(64), activation)
+ self.down1 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
+ norm_layer(128), activation)
+ self.down2 = nn.Sequential(nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
+ norm_layer(256), activation)
+ self.down3 = nn.Sequential(nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),
+ norm_layer(512), activation)
+ if self.deep:
+ self.down4 = nn.Sequential(nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),
+ norm_layer(512), activation)
+
+ ### resnet blocks
+ BN = []
+ for i in range(n_blocks):
+ BN += [
+ ResnetBlock_Adain(512, latent_size=latent_size, padding_type=padding_type, activation=activation)]
+ self.BottleNeck = nn.Sequential(*BN)
+
+ if self.deep:
+ self.up4 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear'),
+ nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(512), activation
+ )
+ self.up3 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear'),
+ nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(256), activation
+ )
+ self.up2 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear'),
+ nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(128), activation
+ )
+ self.up1 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear'),
+ nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(64), activation
+ )
+ self.up0 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear'),
+ nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(32), activation
+ )
+ self.last_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(32, output_nc, kernel_size=7, padding=0),
+ nn.Tanh())
+
+ def forward(self, input, dlatents):
+ x = input # 3*224*224
+
+ skip0 = self.first_layer(x)
+ skip1 = self.down0(skip0)
+ skip2 = self.down1(skip1)
+ skip3 = self.down2(skip2)
+ if self.deep:
+ skip4 = self.down3(skip3)
+ x = self.down4(skip4)
+ else:
+ x = self.down3(skip3)
+
+ for i in range(len(self.BottleNeck)):
+ x = self.BottleNeck[i](x, dlatents)
+
+ if self.deep:
+ x = self.up4(x)
+ x = self.up3(x)
+ x = self.up2(x)
+ x = self.up1(x)
+ x = self.up0(x)
+ x = self.last_layer(x)
+ x = (x + 1) / 2
+
+ return x
+
+class Discriminator(nn.Module):
+ def __init__(self, input_nc, norm_layer=nn.BatchNorm2d, use_sigmoid=False):
+ super(Discriminator, self).__init__()
+
+ kw = 4
+ padw = 1
+ self.down1 = nn.Sequential(
+ nn.Conv2d(input_nc, 64, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)
+ )
+ self.down2 = nn.Sequential(
+ nn.Conv2d(64, 128, kernel_size=kw, stride=2, padding=padw),
+ norm_layer(128), nn.LeakyReLU(0.2, True)
+ )
+ self.down3 = nn.Sequential(
+ nn.Conv2d(128, 256, kernel_size=kw, stride=2, padding=padw),
+ norm_layer(256), nn.LeakyReLU(0.2, True)
+ )
+ self.down4 = nn.Sequential(
+ nn.Conv2d(256, 512, kernel_size=kw, stride=2, padding=padw),
+ norm_layer(512), nn.LeakyReLU(0.2, True)
+ )
+ self.conv1 = nn.Sequential(
+ nn.Conv2d(512, 512, kernel_size=kw, stride=1, padding=padw),
+ norm_layer(512),
+ nn.LeakyReLU(0.2, True)
+ )
+
+ if use_sigmoid:
+ self.conv2 = nn.Sequential(
+ nn.Conv2d(512, 1, kernel_size=kw, stride=1, padding=padw), nn.Sigmoid()
+ )
+ else:
+ self.conv2 = nn.Sequential(
+ nn.Conv2d(512, 1, kernel_size=kw, stride=1, padding=padw)
+ )
+
+ def forward(self, input):
+ out = []
+ x = self.down1(input)
+ out.append(x)
+ x = self.down2(x)
+ out.append(x)
+ x = self.down3(x)
+ out.append(x)
+ x = self.down4(x)
+ out.append(x)
+ x = self.conv1(x)
+ out.append(x)
+ x = self.conv2(x)
+ out.append(x)
+
+ return out
diff --git a/SimSwap/models/fs_networks_fix.py b/SimSwap/models/fs_networks_fix.py
new file mode 100644
index 0000000000000000000000000000000000000000..c7b05250a349d2bb933d652c3b96c8b65f8d07a4
--- /dev/null
+++ b/SimSwap/models/fs_networks_fix.py
@@ -0,0 +1,172 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation. All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import torch
+import torch.nn as nn
+
+
+class InstanceNorm(nn.Module):
+ def __init__(self, epsilon=1e-8):
+ """
+ @notice: avoid in-place ops.
+ https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
+ """
+ super(InstanceNorm, self).__init__()
+ self.epsilon = epsilon
+
+ def forward(self, x):
+ x = x - torch.mean(x, (2, 3), True)
+ tmp = torch.mul(x, x) # or x ** 2
+ tmp = torch.rsqrt(torch.mean(tmp, (2, 3), True) + self.epsilon)
+ return x * tmp
+
+class ApplyStyle(nn.Module):
+ """
+ @ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb
+ """
+ def __init__(self, latent_size, channels):
+ super(ApplyStyle, self).__init__()
+ self.linear = nn.Linear(latent_size, channels * 2)
+
+ def forward(self, x, latent):
+ style = self.linear(latent) # style => [batch_size, n_channels*2]
+ shape = [-1, 2, x.size(1), 1, 1]
+ style = style.view(shape) # [batch_size, 2, n_channels, ...]
+ #x = x * (style[:, 0] + 1.) + style[:, 1]
+ x = x * (style[:, 0] * 1 + 1.) + style[:, 1] * 1
+ return x
+
+class ResnetBlock_Adain(nn.Module):
+ def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True)):
+ super(ResnetBlock_Adain, self).__init__()
+
+ p = 0
+ conv1 = []
+ if padding_type == 'reflect':
+ conv1 += [nn.ReflectionPad2d(1)]
+ elif padding_type == 'replicate':
+ conv1 += [nn.ReplicationPad2d(1)]
+ elif padding_type == 'zero':
+ p = 1
+ else:
+ raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+ conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), InstanceNorm()]
+ self.conv1 = nn.Sequential(*conv1)
+ self.style1 = ApplyStyle(latent_size, dim)
+ self.act1 = activation
+
+ p = 0
+ conv2 = []
+ if padding_type == 'reflect':
+ conv2 += [nn.ReflectionPad2d(1)]
+ elif padding_type == 'replicate':
+ conv2 += [nn.ReplicationPad2d(1)]
+ elif padding_type == 'zero':
+ p = 1
+ else:
+ raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+ conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), InstanceNorm()]
+ self.conv2 = nn.Sequential(*conv2)
+ self.style2 = ApplyStyle(latent_size, dim)
+
+
+ def forward(self, x, dlatents_in_slice):
+ y = self.conv1(x)
+ y = self.style1(y, dlatents_in_slice)
+ y = self.act1(y)
+ y = self.conv2(y)
+ y = self.style2(y, dlatents_in_slice)
+ out = x + y
+ return out
+
+
+
+class Generator_Adain_Upsample(nn.Module):
+ def __init__(self, input_nc, output_nc, latent_size, n_blocks=6, deep=False,
+ norm_layer=nn.BatchNorm2d,
+ padding_type='reflect'):
+ assert (n_blocks >= 0)
+ super(Generator_Adain_Upsample, self).__init__()
+
+ activation = nn.ReLU(True)
+
+ self.deep = deep
+
+ self.first_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(input_nc, 64, kernel_size=7, padding=0),
+ norm_layer(64), activation)
+ ### downsample
+ self.down1 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
+ norm_layer(128), activation)
+ self.down2 = nn.Sequential(nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
+ norm_layer(256), activation)
+ self.down3 = nn.Sequential(nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),
+ norm_layer(512), activation)
+
+ if self.deep:
+ self.down4 = nn.Sequential(nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),
+ norm_layer(512), activation)
+
+ ### resnet blocks
+ BN = []
+ for i in range(n_blocks):
+ BN += [
+ ResnetBlock_Adain(512, latent_size=latent_size, padding_type=padding_type, activation=activation)]
+ self.BottleNeck = nn.Sequential(*BN)
+
+ if self.deep:
+ self.up4 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear',align_corners=False),
+ nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(512), activation
+ )
+ self.up3 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear',align_corners=False),
+ nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(256), activation
+ )
+ self.up2 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear',align_corners=False),
+ nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(128), activation
+ )
+ self.up1 = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='bilinear',align_corners=False),
+ nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
+ nn.BatchNorm2d(64), activation
+ )
+ self.last_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(64, output_nc, kernel_size=7, padding=0))
+
+ def forward(self, input, dlatents):
+ x = input # 3*224*224
+
+ skip1 = self.first_layer(x)
+ skip2 = self.down1(skip1)
+ skip3 = self.down2(skip2)
+ if self.deep:
+ skip4 = self.down3(skip3)
+ x = self.down4(skip4)
+ else:
+ x = self.down3(skip3)
+ bot = []
+ bot.append(x)
+ features = []
+ for i in range(len(self.BottleNeck)):
+ x = self.BottleNeck[i](x, dlatents)
+ bot.append(x)
+
+ if self.deep:
+ x = self.up4(x)
+ features.append(x)
+ x = self.up3(x)
+ features.append(x)
+ x = self.up2(x)
+ features.append(x)
+ x = self.up1(x)
+ features.append(x)
+ x = self.last_layer(x)
+ # x = (x + 1) / 2
+
+ # return x, bot, features, dlatents
+ return x
\ No newline at end of file
diff --git a/SimSwap/models/models.py b/SimSwap/models/models.py
new file mode 100644
index 0000000000000000000000000000000000000000..bf9e84ecaf75220001cdd2151f1102b7385716c6
--- /dev/null
+++ b/SimSwap/models/models.py
@@ -0,0 +1,177 @@
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import Parameter
+from .config import device, num_classes
+
+
+def create_model(opt):
+ #from .pix2pixHD_model import Pix2PixHDModel, InferenceModel
+ from .fs_model import fsModel
+ model = fsModel()
+
+ model.initialize(opt)
+ if opt.verbose:
+ print("model [%s] was created" % (model.name()))
+
+ if opt.isTrain and len(opt.gpu_ids) and not opt.fp16:
+ model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids)
+
+ return model
+
+
+
+class SEBlock(nn.Module):
+ def __init__(self, channel, reduction=16):
+ super(SEBlock, self).__init__()
+ self.avg_pool = nn.AdaptiveAvgPool2d(1)
+ self.fc = nn.Sequential(
+ nn.Linear(channel, channel // reduction),
+ nn.PReLU(),
+ nn.Linear(channel // reduction, channel),
+ nn.Sigmoid()
+ )
+
+ def forward(self, x):
+ b, c, _, _ = x.size()
+ y = self.avg_pool(x).view(b, c)
+ y = self.fc(y).view(b, c, 1, 1)
+ return x * y
+
+
+class IRBlock(nn.Module):
+ expansion = 1
+
+ def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True):
+ super(IRBlock, self).__init__()
+ self.bn0 = nn.BatchNorm2d(inplanes)
+ self.conv1 = conv3x3(inplanes, inplanes)
+ self.bn1 = nn.BatchNorm2d(inplanes)
+ self.prelu = nn.PReLU()
+ self.conv2 = conv3x3(inplanes, planes, stride)
+ self.bn2 = nn.BatchNorm2d(planes)
+ self.downsample = downsample
+ self.stride = stride
+ self.use_se = use_se
+ if self.use_se:
+ self.se = SEBlock(planes)
+
+ def forward(self, x):
+ residual = x
+ out = self.bn0(x)
+ out = self.conv1(out)
+ out = self.bn1(out)
+ out = self.prelu(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ if self.use_se:
+ out = self.se(out)
+
+ if self.downsample is not None:
+ residual = self.downsample(x)
+
+ out += residual
+ out = self.prelu(out)
+
+ return out
+
+
+class ResNet(nn.Module):
+
+ def __init__(self, block, layers, use_se=True):
+ self.inplanes = 64
+ self.use_se = use_se
+ super(ResNet, self).__init__()
+ self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(64)
+ self.prelu = nn.PReLU()
+ self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
+ self.layer1 = self._make_layer(block, 64, layers[0])
+ self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+ self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+ self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+ self.bn2 = nn.BatchNorm2d(512)
+ self.dropout = nn.Dropout()
+ self.fc = nn.Linear(512 * 7 * 7, 512)
+ self.bn3 = nn.BatchNorm1d(512)
+
+ for m in self.modules():
+ if isinstance(m, nn.Conv2d):
+ nn.init.xavier_normal_(m.weight)
+ elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+ elif isinstance(m, nn.Linear):
+ nn.init.xavier_normal_(m.weight)
+ nn.init.constant_(m.bias, 0)
+
+ def _make_layer(self, block, planes, blocks, stride=1):
+ downsample = None
+ if stride != 1 or self.inplanes != planes * block.expansion:
+ downsample = nn.Sequential(
+ nn.Conv2d(self.inplanes, planes * block.expansion,
+ kernel_size=1, stride=stride, bias=False),
+ nn.BatchNorm2d(planes * block.expansion),
+ )
+
+ layers = []
+ layers.append(block(self.inplanes, planes, stride, downsample, use_se=self.use_se))
+ self.inplanes = planes
+ for i in range(1, blocks):
+ layers.append(block(self.inplanes, planes, use_se=self.use_se))
+
+ return nn.Sequential(*layers)
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.bn1(x)
+ x = self.prelu(x)
+ x = self.maxpool(x)
+
+ x = self.layer1(x)
+ x = self.layer2(x)
+ x = self.layer3(x)
+ x = self.layer4(x)
+
+ x = self.bn2(x)
+ x = self.dropout(x)
+ x = x.view(x.size(0), -1)
+ x = self.fc(x)
+ x = self.bn3(x)
+
+ return x
+
+
+class ArcMarginModel(nn.Module):
+ def __init__(self, args):
+ super(ArcMarginModel, self).__init__()
+
+ self.weight = Parameter(torch.FloatTensor(num_classes, args.emb_size))
+ nn.init.xavier_uniform_(self.weight)
+
+ self.easy_margin = args.easy_margin
+ self.m = args.margin_m
+ self.s = args.margin_s
+
+ self.cos_m = math.cos(self.m)
+ self.sin_m = math.sin(self.m)
+ self.th = math.cos(math.pi - self.m)
+ self.mm = math.sin(math.pi - self.m) * self.m
+
+ def forward(self, input, label):
+ x = F.normalize(input)
+ W = F.normalize(self.weight)
+ cosine = F.linear(x, W)
+ sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
+ phi = cosine * self.cos_m - sine * self.sin_m # cos(theta + m)
+ if self.easy_margin:
+ phi = torch.where(cosine > 0, phi, cosine)
+ else:
+ phi = torch.where(cosine > self.th, phi, cosine - self.mm)
+ one_hot = torch.zeros(cosine.size(), device=device)
+ one_hot.scatter_(1, label.view(-1, 1).long(), 1)
+ output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
+ output *= self.s
+ return output
diff --git a/SimSwap/models/networks.py b/SimSwap/models/networks.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c7eb4de0421ccb62e75bc47b06accb026f1767f
--- /dev/null
+++ b/SimSwap/models/networks.py
@@ -0,0 +1,845 @@
+import torch
+import torch.nn as nn
+import functools
+from torch.autograd import Variable
+import numpy as np
+from torchvision import transforms
+import torch.nn.functional as F
+
+###############################################################################
+# Functions
+###############################################################################
+def weights_init(m):
+ classname = m.__class__.__name__
+ if classname.find('Conv') != -1:
+ m.weight.data.normal_(0.0, 0.02)
+ elif classname.find('BatchNorm2d') != -1:
+ m.weight.data.normal_(1.0, 0.02)
+ m.bias.data.fill_(0)
+
+def get_norm_layer(norm_type='instance'):
+ if norm_type == 'batch':
+ norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
+ elif norm_type == 'instance':
+ norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)
+ else:
+ raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+ return norm_layer
+
+def define_G(input_nc, output_nc, ngf, netG, n_downsample_global=3, n_blocks_global=9, n_local_enhancers=1,
+ n_blocks_local=3, norm='instance', gpu_ids=[]):
+ norm_layer = get_norm_layer(norm_type=norm)
+ if netG == 'global':
+ netG = GlobalGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer)
+ elif netG == 'local':
+ netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global,
+ n_local_enhancers, n_blocks_local, norm_layer)
+ elif netG == 'encoder':
+ netG = Encoder(input_nc, output_nc, ngf, n_downsample_global, norm_layer)
+ else:
+ raise('generator not implemented!')
+ print(netG)
+ if len(gpu_ids) > 0:
+ assert(torch.cuda.is_available())
+ netG.cuda(gpu_ids[0])
+ netG.apply(weights_init)
+ return netG
+
+def define_G_Adain(input_nc, output_nc, latent_size, ngf, netG, n_downsample_global=2, n_blocks_global=4, norm='instance', gpu_ids=[]):
+ norm_layer = get_norm_layer(norm_type=norm)
+ netG = Generator_Adain(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
+ print(netG)
+ if len(gpu_ids) > 0:
+ assert(torch.cuda.is_available())
+ netG.cuda(gpu_ids[0])
+ netG.apply(weights_init)
+ return netG
+
+def define_G_Adain_Mask(input_nc, output_nc, latent_size, ngf, netG, n_downsample_global=2, n_blocks_global=4, norm='instance', gpu_ids=[]):
+ norm_layer = get_norm_layer(norm_type=norm)
+ netG = Generator_Adain_Mask(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
+ print(netG)
+ if len(gpu_ids) > 0:
+ assert(torch.cuda.is_available())
+ netG.cuda(gpu_ids[0])
+ netG.apply(weights_init)
+ return netG
+
+def define_G_Adain_Upsample(input_nc, output_nc, latent_size, ngf, netG, n_downsample_global=2, n_blocks_global=4, norm='instance', gpu_ids=[]):
+ norm_layer = get_norm_layer(norm_type=norm)
+ netG = Generator_Adain_Upsample(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
+ print(netG)
+ if len(gpu_ids) > 0:
+ assert(torch.cuda.is_available())
+ netG.cuda(gpu_ids[0])
+ netG.apply(weights_init)
+ return netG
+
+def define_G_Adain_2(input_nc, output_nc, latent_size, ngf, netG, n_downsample_global=2, n_blocks_global=4, norm='instance', gpu_ids=[]):
+ norm_layer = get_norm_layer(norm_type=norm)
+ netG = Generator_Adain_2(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
+ print(netG)
+ if len(gpu_ids) > 0:
+ assert(torch.cuda.is_available())
+ netG.cuda(gpu_ids[0])
+ netG.apply(weights_init)
+ return netG
+
+def define_D(input_nc, ndf, n_layers_D, norm='instance', use_sigmoid=False, num_D=1, getIntermFeat=False, gpu_ids=[]):
+ norm_layer = get_norm_layer(norm_type=norm)
+ netD = MultiscaleDiscriminator(input_nc, ndf, n_layers_D, norm_layer, use_sigmoid, num_D, getIntermFeat)
+ print(netD)
+ if len(gpu_ids) > 0:
+ assert(torch.cuda.is_available())
+ netD.cuda(gpu_ids[0])
+ netD.apply(weights_init)
+ return netD
+
+def print_network(net):
+ if isinstance(net, list):
+ net = net[0]
+ num_params = 0
+ for param in net.parameters():
+ num_params += param.numel()
+ print(net)
+ print('Total number of parameters: %d' % num_params)
+
+##############################################################################
+# Losses
+##############################################################################
+class GANLoss(nn.Module):
+ def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0,
+ tensor=torch.FloatTensor, opt=None):
+ super(GANLoss, self).__init__()
+ self.real_label = target_real_label
+ self.fake_label = target_fake_label
+ self.real_label_tensor = None
+ self.fake_label_tensor = None
+ self.zero_tensor = None
+ self.Tensor = tensor
+ self.gan_mode = gan_mode
+ self.opt = opt
+ if gan_mode == 'ls':
+ pass
+ elif gan_mode == 'original':
+ pass
+ elif gan_mode == 'w':
+ pass
+ elif gan_mode == 'hinge':
+ pass
+ else:
+ raise ValueError('Unexpected gan_mode {}'.format(gan_mode))
+
+ def get_target_tensor(self, input, target_is_real):
+ if target_is_real:
+ if self.real_label_tensor is None:
+ self.real_label_tensor = self.Tensor(1).fill_(self.real_label)
+ self.real_label_tensor.requires_grad_(False)
+ return self.real_label_tensor.expand_as(input)
+ else:
+ if self.fake_label_tensor is None:
+ self.fake_label_tensor = self.Tensor(1).fill_(self.fake_label)
+ self.fake_label_tensor.requires_grad_(False)
+ return self.fake_label_tensor.expand_as(input)
+
+ def get_zero_tensor(self, input):
+ if self.zero_tensor is None:
+ self.zero_tensor = self.Tensor(1).fill_(0)
+ self.zero_tensor.requires_grad_(False)
+ return self.zero_tensor.expand_as(input)
+
+ def loss(self, input, target_is_real, for_discriminator=True):
+ if self.gan_mode == 'original': # cross entropy loss
+ target_tensor = self.get_target_tensor(input, target_is_real)
+ loss = F.binary_cross_entropy_with_logits(input, target_tensor)
+ return loss
+ elif self.gan_mode == 'ls':
+ target_tensor = self.get_target_tensor(input, target_is_real)
+ return F.mse_loss(input, target_tensor)
+ elif self.gan_mode == 'hinge':
+ if for_discriminator:
+ if target_is_real:
+ minval = torch.min(input - 1, self.get_zero_tensor(input))
+ loss = -torch.mean(minval)
+ else:
+ minval = torch.min(-input - 1, self.get_zero_tensor(input))
+ loss = -torch.mean(minval)
+ else:
+ assert target_is_real, "The generator's hinge loss must be aiming for real"
+ loss = -torch.mean(input)
+ return loss
+ else:
+ # wgan
+ if target_is_real:
+ return -input.mean()
+ else:
+ return input.mean()
+
+ def __call__(self, input, target_is_real, for_discriminator=True):
+ # computing loss is a bit complicated because |input| may not be
+ # a tensor, but list of tensors in case of multiscale discriminator
+ if isinstance(input, list):
+ loss = 0
+ for pred_i in input:
+ if isinstance(pred_i, list):
+ pred_i = pred_i[-1]
+ loss_tensor = self.loss(pred_i, target_is_real, for_discriminator)
+ bs = 1 if len(loss_tensor.size()) == 0 else loss_tensor.size(0)
+ new_loss = torch.mean(loss_tensor.view(bs, -1), dim=1)
+ loss += new_loss
+ return loss / len(input)
+ else:
+ return self.loss(input, target_is_real, for_discriminator)
+
+class VGGLoss(nn.Module):
+ def __init__(self, gpu_ids):
+ super(VGGLoss, self).__init__()
+ self.vgg = Vgg19().cuda()
+ self.criterion = nn.L1Loss()
+ self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]
+
+ def forward(self, x, y):
+ x_vgg, y_vgg = self.vgg(x), self.vgg(y)
+ loss = 0
+ for i in range(len(x_vgg)):
+ loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
+ return loss
+
+##############################################################################
+# Generator
+##############################################################################
+class LocalEnhancer(nn.Module):
+ def __init__(self, input_nc, output_nc, ngf=32, n_downsample_global=3, n_blocks_global=9,
+ n_local_enhancers=1, n_blocks_local=3, norm_layer=nn.BatchNorm2d, padding_type='reflect'):
+ super(LocalEnhancer, self).__init__()
+ self.n_local_enhancers = n_local_enhancers
+
+ ###### global generator model #####
+ ngf_global = ngf * (2**n_local_enhancers)
+ model_global = GlobalGenerator(input_nc, output_nc, ngf_global, n_downsample_global, n_blocks_global, norm_layer).model
+ model_global = [model_global[i] for i in range(len(model_global)-3)] # get rid of final convolution layers
+ self.model = nn.Sequential(*model_global)
+
+ ###### local enhancer layers #####
+ for n in range(1, n_local_enhancers+1):
+ ### downsample
+ ngf_global = ngf * (2**(n_local_enhancers-n))
+ model_downsample = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0),
+ norm_layer(ngf_global), nn.ReLU(True),
+ nn.Conv2d(ngf_global, ngf_global * 2, kernel_size=3, stride=2, padding=1),
+ norm_layer(ngf_global * 2), nn.ReLU(True)]
+ ### residual blocks
+ model_upsample = []
+ for i in range(n_blocks_local):
+ model_upsample += [ResnetBlock(ngf_global * 2, padding_type=padding_type, norm_layer=norm_layer)]
+
+ ### upsample
+ model_upsample += [nn.ConvTranspose2d(ngf_global * 2, ngf_global, kernel_size=3, stride=2, padding=1, output_padding=1),
+ norm_layer(ngf_global), nn.ReLU(True)]
+
+ ### final convolution
+ if n == n_local_enhancers:
+ model_upsample += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
+
+ setattr(self, 'model'+str(n)+'_1', nn.Sequential(*model_downsample))
+ setattr(self, 'model'+str(n)+'_2', nn.Sequential(*model_upsample))
+
+ self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
+
+ def forward(self, input):
+ ### create input pyramid
+ input_downsampled = [input]
+ for i in range(self.n_local_enhancers):
+ input_downsampled.append(self.downsample(input_downsampled[-1]))
+
+ ### output at coarest level
+ output_prev = self.model(input_downsampled[-1])
+ ### build up one layer at a time
+ for n_local_enhancers in range(1, self.n_local_enhancers+1):
+ model_downsample = getattr(self, 'model'+str(n_local_enhancers)+'_1')
+ model_upsample = getattr(self, 'model'+str(n_local_enhancers)+'_2')
+ input_i = input_downsampled[self.n_local_enhancers-n_local_enhancers]
+ output_prev = model_upsample(model_downsample(input_i) + output_prev)
+ return output_prev
+
+class GlobalGenerator(nn.Module):
+ def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=nn.BatchNorm2d,
+ padding_type='reflect'):
+ assert(n_blocks >= 0)
+ super(GlobalGenerator, self).__init__()
+ activation = nn.ReLU(True)
+
+ model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
+ ### downsample
+ for i in range(n_downsampling):
+ mult = 2**i
+ model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
+ norm_layer(ngf * mult * 2), activation]
+
+ ### resnet blocks
+ mult = 2**n_downsampling
+ for i in range(n_blocks):
+ model += [ResnetBlock(ngf * mult, padding_type=padding_type, activation=activation, norm_layer=norm_layer)]
+
+ ### upsample
+ for i in range(n_downsampling):
+ mult = 2**(n_downsampling - i)
+ model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1),
+ norm_layer(int(ngf * mult / 2)), activation]
+ model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
+ self.model = nn.Sequential(*model)
+
+ def forward(self, input):
+ return self.model(input)
+
+# Define a resnet block
+class ResnetBlock(nn.Module):
+ def __init__(self, dim, padding_type, norm_layer, activation=nn.ReLU(True), use_dropout=False):
+ super(ResnetBlock, self).__init__()
+ self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, activation, use_dropout)
+
+ def build_conv_block(self, dim, padding_type, norm_layer, activation, use_dropout):
+ conv_block = []
+ p = 0
+ if padding_type == 'reflect':
+ conv_block += [nn.ReflectionPad2d(1)]
+ elif padding_type == 'replicate':
+ conv_block += [nn.ReplicationPad2d(1)]
+ elif padding_type == 'zero':
+ p = 1
+ else:
+ raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+ conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p),
+ norm_layer(dim),
+ activation]
+ if use_dropout:
+ conv_block += [nn.Dropout(0.5)]
+
+ p = 0
+ if padding_type == 'reflect':
+ conv_block += [nn.ReflectionPad2d(1)]
+ elif padding_type == 'replicate':
+ conv_block += [nn.ReplicationPad2d(1)]
+ elif padding_type == 'zero':
+ p = 1
+ else:
+ raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+ conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p),
+ norm_layer(dim)]
+
+ return nn.Sequential(*conv_block)
+
+ def forward(self, x):
+ out = x + self.conv_block(x)
+ return out
+
+class InstanceNorm(nn.Module):
+ def __init__(self, epsilon=1e-8):
+ """
+ @notice: avoid in-place ops.
+ https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
+ """
+ super(InstanceNorm, self).__init__()
+ self.epsilon = epsilon
+
+ def forward(self, x):
+ x = x - torch.mean(x, (2, 3), True)
+ tmp = torch.mul(x, x) # or x ** 2
+ tmp = torch.rsqrt(torch.mean(tmp, (2, 3), True) + self.epsilon)
+ return x * tmp
+
+class SpecificNorm(nn.Module):
+ def __init__(self, epsilon=1e-8):
+ """
+ @notice: avoid in-place ops.
+ https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
+ """
+ super(SpecificNorm, self).__init__()
+ self.mean = np.array([0.485, 0.456, 0.406])
+ self.mean = torch.from_numpy(self.mean).float().cuda()
+ self.mean = self.mean.view([1, 3, 1, 1])
+
+ self.std = np.array([0.229, 0.224, 0.225])
+ self.std = torch.from_numpy(self.std).float().cuda()
+ self.std = self.std.view([1, 3, 1, 1])
+
+ def forward(self, x):
+ mean = self.mean.expand([1, 3, x.shape[2], x.shape[3]])
+ std = self.std.expand([1, 3, x.shape[2], x.shape[3]])
+
+ x = (x - mean) / std
+
+ return x
+
+class ApplyStyle(nn.Module):
+ """
+ @ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb
+ """
+ def __init__(self, latent_size, channels):
+ super(ApplyStyle, self).__init__()
+ self.linear = nn.Linear(latent_size, channels * 2)
+
+ def forward(self, x, latent):
+ style = self.linear(latent) # style => [batch_size, n_channels*2]
+ shape = [-1, 2, x.size(1), 1, 1]
+ style = style.view(shape) # [batch_size, 2, n_channels, ...]
+ x = x * (style[:, 0] + 1.) + style[:, 1]
+ return x
+
+class ResnetBlock_Adain(nn.Module):
+ def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True)):
+ super(ResnetBlock_Adain, self).__init__()
+
+ p = 0
+ conv1 = []
+ if padding_type == 'reflect':
+ conv1 += [nn.ReflectionPad2d(1)]
+ elif padding_type == 'replicate':
+ conv1 += [nn.ReplicationPad2d(1)]
+ elif padding_type == 'zero':
+ p = 1
+ else:
+ raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+ conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), InstanceNorm()]
+ self.conv1 = nn.Sequential(*conv1)
+ self.style1 = ApplyStyle(latent_size, dim)
+ self.act1 = activation
+
+ p = 0
+ conv2 = []
+ if padding_type == 'reflect':
+ conv2 += [nn.ReflectionPad2d(1)]
+ elif padding_type == 'replicate':
+ conv2 += [nn.ReplicationPad2d(1)]
+ elif padding_type == 'zero':
+ p = 1
+ else:
+ raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+ conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), InstanceNorm()]
+ self.conv2 = nn.Sequential(*conv2)
+ self.style2 = ApplyStyle(latent_size, dim)
+
+
+ def forward(self, x, dlatents_in_slice):
+ y = self.conv1(x)
+ y = self.style1(y, dlatents_in_slice)
+ y = self.act1(y)
+ y = self.conv2(y)
+ y = self.style2(y, dlatents_in_slice)
+ out = x + y
+ return out
+
+class UpBlock_Adain(nn.Module):
+ def __init__(self, dim_in, dim_out, latent_size, padding_type, activation=nn.ReLU(True)):
+ super(UpBlock_Adain, self).__init__()
+
+ p = 0
+ conv1 = [nn.Upsample(scale_factor=2, mode='bilinear')]
+ if padding_type == 'reflect':
+ conv1 += [nn.ReflectionPad2d(1)]
+ elif padding_type == 'replicate':
+ conv1 += [nn.ReplicationPad2d(1)]
+ elif padding_type == 'zero':
+ p = 1
+ else:
+ raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+ conv1 += [nn.Conv2d(dim_in, dim_out, kernel_size=3, padding = p), InstanceNorm()]
+ self.conv1 = nn.Sequential(*conv1)
+ self.style1 = ApplyStyle(latent_size, dim_out)
+ self.act1 = activation
+
+
+ def forward(self, x, dlatents_in_slice):
+ y = self.conv1(x)
+ y = self.style1(y, dlatents_in_slice)
+ y = self.act1(y)
+ return y
+
+class Encoder(nn.Module):
+ def __init__(self, input_nc, output_nc, ngf=32, n_downsampling=4, norm_layer=nn.BatchNorm2d):
+ super(Encoder, self).__init__()
+ self.output_nc = output_nc
+
+ model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0),
+ norm_layer(ngf), nn.ReLU(True)]
+ ### downsample
+ for i in range(n_downsampling):
+ mult = 2**i
+ model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
+ norm_layer(ngf * mult * 2), nn.ReLU(True)]
+
+ ### upsample
+ for i in range(n_downsampling):
+ mult = 2**(n_downsampling - i)
+ model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1),
+ norm_layer(int(ngf * mult / 2)), nn.ReLU(True)]
+
+ model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
+ self.model = nn.Sequential(*model)
+
+ def forward(self, input, inst):
+ outputs = self.model(input)
+
+ # instance-wise average pooling
+ outputs_mean = outputs.clone()
+ inst_list = np.unique(inst.cpu().numpy().astype(int))
+ for i in inst_list:
+ for b in range(input.size()[0]):
+ indices = (inst[b:b+1] == int(i)).nonzero() # n x 4
+ for j in range(self.output_nc):
+ output_ins = outputs[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]]
+ mean_feat = torch.mean(output_ins).expand_as(output_ins)
+ outputs_mean[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]] = mean_feat
+ return outputs_mean
+
+
+class Generator_Adain(nn.Module):
+ def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
+ padding_type='reflect'):
+ assert (n_blocks >= 0)
+ super(Generator_Adain, self).__init__()
+ activation = nn.ReLU(True)
+
+ Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
+ ### downsample
+ for i in range(n_downsampling):
+ mult = 2 ** i
+ Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
+ norm_layer(ngf * mult * 2), activation]
+ self.Encoder = nn.Sequential(*Enc)
+
+ ### resnet blocks
+ BN = []
+ mult = 2 ** n_downsampling
+ for i in range(n_blocks):
+ BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
+ self.BottleNeck = nn.Sequential(*BN)
+ '''self.ResBlockAdain1 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
+ activation=activation)
+ self.ResBlockAdain2 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
+ activation=activation)
+ self.ResBlockAdain3 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
+ activation=activation)
+ self.ResBlockAdain4 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
+ activation=activation)'''
+
+ ### upsample
+ Dec = []
+ for i in range(n_downsampling):
+ mult = 2 ** (n_downsampling - i)
+ Dec += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1,
+ output_padding=1),
+ norm_layer(int(ngf * mult / 2)), activation]
+ Dec += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
+
+ self.Decoder = nn.Sequential(*Dec)
+ #self.model = nn.Sequential(*model)
+ self.spNorm = SpecificNorm()
+
+ def forward(self, input, dlatents):
+ x = input
+ x = self.Encoder(x)
+
+
+ for i in range(len(self.BottleNeck)):
+ x = self.BottleNeck[i](x, dlatents)
+ '''x = self.ResBlockAdain1(x, dlatents)
+ x = self.ResBlockAdain2(x, dlatents)
+ x = self.ResBlockAdain3(x, dlatents)
+ x = self.ResBlockAdain4(x, dlatents)'''
+
+ x = self.Decoder(x)
+
+ x = (x + 1) / 2
+ x = self.spNorm(x)
+
+ return x
+
+class Generator_Adain_Mask(nn.Module):
+ def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
+ padding_type='reflect'):
+ assert (n_blocks >= 0)
+ super(Generator_Adain_Mask, self).__init__()
+ activation = nn.ReLU(True)
+
+ Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
+ ### downsample
+ for i in range(n_downsampling):
+ mult = 2 ** i
+ Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
+ norm_layer(ngf * mult * 2), activation]
+ self.Encoder = nn.Sequential(*Enc)
+
+ ### resnet blocks
+ BN = []
+ mult = 2 ** n_downsampling
+ for i in range(n_blocks):
+ BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
+ self.BottleNeck = nn.Sequential(*BN)
+
+ ### upsample
+ Dec = []
+ for i in range(n_downsampling):
+ mult = 2 ** (n_downsampling - i)
+ Dec += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1,
+ output_padding=1),
+ norm_layer(int(ngf * mult / 2)), activation]
+ Fake_out = [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
+ Mast_out = [nn.ReflectionPad2d(3), nn.Conv2d(ngf, 1, kernel_size=7, padding=0), nn.Sigmoid()]
+
+ self.Decoder = nn.Sequential(*Dec)
+ #self.model = nn.Sequential(*model)
+ self.spNorm = SpecificNorm()
+ self.Fake_out = nn.Sequential(*Fake_out)
+ self.Mask_out = nn.Sequential(*Mast_out)
+
+ def forward(self, input, dlatents):
+ x = input
+ x = self.Encoder(x)
+
+
+ for i in range(len(self.BottleNeck)):
+ x = self.BottleNeck[i](x, dlatents)
+
+ x = self.Decoder(x)
+
+ fake_out = self.Fake_out(x)
+ mask_out = self.Mask_out(x)
+
+ fake_out = (fake_out + 1) / 2
+ fake_out = self.spNorm(fake_out)
+
+ generated = fake_out * mask_out + input * (1-mask_out)
+
+ return generated, mask_out
+
+class Generator_Adain_Upsample(nn.Module):
+ def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
+ padding_type='reflect'):
+ assert (n_blocks >= 0)
+ super(Generator_Adain_Upsample, self).__init__()
+ activation = nn.ReLU(True)
+
+ Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
+ ### downsample
+ for i in range(n_downsampling):
+ mult = 2 ** i
+ Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
+ norm_layer(ngf * mult * 2), activation]
+ self.Encoder = nn.Sequential(*Enc)
+
+ ### resnet blocks
+ BN = []
+ mult = 2 ** n_downsampling
+ for i in range(n_blocks):
+ BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
+ self.BottleNeck = nn.Sequential(*BN)
+
+ ### upsample
+ Dec = []
+ for i in range(n_downsampling):
+ mult = 2 ** (n_downsampling - i)
+ '''Dec += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1,
+ output_padding=1),
+ norm_layer(int(ngf * mult / 2)), activation]'''
+ Dec += [nn.Upsample(scale_factor=2, mode='bilinear'),
+ nn.Conv2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=1, padding=1),
+ norm_layer(int(ngf * mult / 2)), activation]
+ Dec += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
+
+ self.Decoder = nn.Sequential(*Dec)
+ self.spNorm = SpecificNorm()
+
+ def forward(self, input, dlatents):
+ x = input
+ x = self.Encoder(x)
+
+
+ for i in range(len(self.BottleNeck)):
+ x = self.BottleNeck[i](x, dlatents)
+
+ x = self.Decoder(x)
+
+ x = (x + 1) / 2
+ x = self.spNorm(x)
+
+ return x
+
+class Generator_Adain_2(nn.Module):
+ def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
+ padding_type='reflect'):
+ assert (n_blocks >= 0)
+ super(Generator_Adain_2, self).__init__()
+ activation = nn.ReLU(True)
+
+ Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
+ ### downsample
+ for i in range(n_downsampling):
+ mult = 2 ** i
+ Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
+ norm_layer(ngf * mult * 2), activation]
+ self.Encoder = nn.Sequential(*Enc)
+
+ ### resnet blocks
+ BN = []
+ mult = 2 ** n_downsampling
+ for i in range(n_blocks):
+ BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
+ self.BottleNeck = nn.Sequential(*BN)
+
+ ### upsample
+ Dec = []
+ for i in range(n_downsampling):
+ mult = 2 ** (n_downsampling - i)
+ Dec += [UpBlock_Adain(dim_in=ngf * mult, dim_out=int(ngf * mult / 2), latent_size=latent_size, padding_type=padding_type)]
+ layer_out = [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
+
+ self.Decoder = nn.Sequential(*Dec)
+ #self.model = nn.Sequential(*model)
+ self.spNorm = SpecificNorm()
+ self.layer_out = nn.Sequential(*layer_out)
+
+ def forward(self, input, dlatents):
+ x = input
+ x = self.Encoder(x)
+
+
+ for i in range(len(self.BottleNeck)):
+ x = self.BottleNeck[i](x, dlatents)
+
+ for i in range(len(self.Decoder)):
+ x = self.Decoder[i](x, dlatents)
+
+ x = self.layer_out(x)
+
+ x = (x + 1) / 2
+ x = self.spNorm(x)
+
+ return x
+
+class MultiscaleDiscriminator(nn.Module):
+ def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d,
+ use_sigmoid=False, num_D=3, getIntermFeat=False):
+ super(MultiscaleDiscriminator, self).__init__()
+ self.num_D = num_D
+ self.n_layers = n_layers
+ self.getIntermFeat = getIntermFeat
+
+ for i in range(num_D):
+ netD = NLayerDiscriminator(input_nc, ndf, n_layers, norm_layer, use_sigmoid, getIntermFeat)
+ if getIntermFeat:
+ for j in range(n_layers+2):
+ setattr(self, 'scale'+str(i)+'_layer'+str(j), getattr(netD, 'model'+str(j)))
+ else:
+ setattr(self, 'layer'+str(i), netD.model)
+
+ self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
+
+ def singleD_forward(self, model, input):
+ if self.getIntermFeat:
+ result = [input]
+ for i in range(len(model)):
+ result.append(model[i](result[-1]))
+ return result[1:]
+ else:
+ return [model(input)]
+
+ def forward(self, input):
+ num_D = self.num_D
+ result = []
+ input_downsampled = input
+ for i in range(num_D):
+ if self.getIntermFeat:
+ model = [getattr(self, 'scale'+str(num_D-1-i)+'_layer'+str(j)) for j in range(self.n_layers+2)]
+ else:
+ model = getattr(self, 'layer'+str(num_D-1-i))
+ result.append(self.singleD_forward(model, input_downsampled))
+ if i != (num_D-1):
+ input_downsampled = self.downsample(input_downsampled)
+ return result
+
+# Defines the PatchGAN discriminator with the specified arguments.
+class NLayerDiscriminator(nn.Module):
+ def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False, getIntermFeat=False):
+ super(NLayerDiscriminator, self).__init__()
+ self.getIntermFeat = getIntermFeat
+ self.n_layers = n_layers
+
+ kw = 4
+ padw = 1
+ sequence = [[nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]]
+
+ nf = ndf
+ for n in range(1, n_layers):
+ nf_prev = nf
+ nf = min(nf * 2, 512)
+ sequence += [[
+ nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=2, padding=padw),
+ norm_layer(nf), nn.LeakyReLU(0.2, True)
+ ]]
+
+ nf_prev = nf
+ nf = min(nf * 2, 512)
+ sequence += [[
+ nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw),
+ norm_layer(nf),
+ nn.LeakyReLU(0.2, True)
+ ]]
+
+ if use_sigmoid:
+ sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw), nn.Sigmoid()]]
+ else:
+ sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]]
+
+ if getIntermFeat:
+ for n in range(len(sequence)):
+ setattr(self, 'model'+str(n), nn.Sequential(*sequence[n]))
+ else:
+ sequence_stream = []
+ for n in range(len(sequence)):
+ sequence_stream += sequence[n]
+ self.model = nn.Sequential(*sequence_stream)
+
+ def forward(self, input):
+ if self.getIntermFeat:
+ res = [input]
+ for n in range(self.n_layers+2):
+ model = getattr(self, 'model'+str(n))
+ res.append(model(res[-1]))
+ return res[1:]
+ else:
+ return self.model(input)
+
+from torchvision import models
+class Vgg19(torch.nn.Module):
+ def __init__(self, requires_grad=False):
+ super(Vgg19, self).__init__()
+ vgg_pretrained_features = models.vgg19(pretrained=True).features
+ self.slice1 = torch.nn.Sequential()
+ self.slice2 = torch.nn.Sequential()
+ self.slice3 = torch.nn.Sequential()
+ self.slice4 = torch.nn.Sequential()
+ self.slice5 = torch.nn.Sequential()
+ for x in range(2):
+ self.slice1.add_module(str(x), vgg_pretrained_features[x])
+ for x in range(2, 7):
+ self.slice2.add_module(str(x), vgg_pretrained_features[x])
+ for x in range(7, 12):
+ self.slice3.add_module(str(x), vgg_pretrained_features[x])
+ for x in range(12, 21):
+ self.slice4.add_module(str(x), vgg_pretrained_features[x])
+ for x in range(21, 30):
+ self.slice5.add_module(str(x), vgg_pretrained_features[x])
+ if not requires_grad:
+ for param in self.parameters():
+ param.requires_grad = False
+
+ def forward(self, X):
+ h_relu1 = self.slice1(X)
+ h_relu2 = self.slice2(h_relu1)
+ h_relu3 = self.slice3(h_relu2)
+ h_relu4 = self.slice4(h_relu3)
+ h_relu5 = self.slice5(h_relu4)
+ out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
+ return out
diff --git a/SimSwap/models/pix2pixHD_model.py b/SimSwap/models/pix2pixHD_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..fafdec0b7f35499070ef9eb015a72584b494164d
--- /dev/null
+++ b/SimSwap/models/pix2pixHD_model.py
@@ -0,0 +1,304 @@
+import numpy as np
+import torch
+import os
+from torch.autograd import Variable
+from util.image_pool import ImagePool
+from .base_model import BaseModel
+from . import networks
+
+class Pix2PixHDModel(BaseModel):
+ def name(self):
+ return 'Pix2PixHDModel'
+
+ def init_loss_filter(self, use_gan_feat_loss, use_vgg_loss):
+ flags = (True, use_gan_feat_loss, use_vgg_loss, True, True)
+ def loss_filter(g_gan, g_gan_feat, g_vgg, d_real, d_fake):
+ return [l for (l,f) in zip((g_gan,g_gan_feat,g_vgg,d_real,d_fake),flags) if f]
+ return loss_filter
+
+ def initialize(self, opt):
+ BaseModel.initialize(self, opt)
+ if opt.resize_or_crop != 'none' or not opt.isTrain: # when training at full res this causes OOM
+ torch.backends.cudnn.benchmark = True
+ self.isTrain = opt.isTrain
+ self.use_features = opt.instance_feat or opt.label_feat
+ self.gen_features = self.use_features and not self.opt.load_features
+ input_nc = opt.label_nc if opt.label_nc != 0 else opt.input_nc
+
+ ##### define networks
+ # Generator network
+ netG_input_nc = input_nc
+ if not opt.no_instance:
+ netG_input_nc += 1
+ if self.use_features:
+ netG_input_nc += opt.feat_num
+ self.netG = networks.define_G(netG_input_nc, opt.output_nc, opt.ngf, opt.netG,
+ opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers,
+ opt.n_blocks_local, opt.norm, gpu_ids=self.gpu_ids)
+
+ # Discriminator network
+ if self.isTrain:
+ use_sigmoid = opt.no_lsgan
+ netD_input_nc = input_nc + opt.output_nc
+ if not opt.no_instance:
+ netD_input_nc += 1
+ self.netD = networks.define_D(netD_input_nc, opt.ndf, opt.n_layers_D, opt.norm, use_sigmoid,
+ opt.num_D, not opt.no_ganFeat_loss, gpu_ids=self.gpu_ids)
+
+ ### Encoder network
+ if self.gen_features:
+ self.netE = networks.define_G(opt.output_nc, opt.feat_num, opt.nef, 'encoder',
+ opt.n_downsample_E, norm=opt.norm, gpu_ids=self.gpu_ids)
+ if self.opt.verbose:
+ print('---------- Networks initialized -------------')
+
+ # load networks
+ if not self.isTrain or opt.continue_train or opt.load_pretrain:
+ pretrained_path = '' if not self.isTrain else opt.load_pretrain
+ self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)
+ if self.isTrain:
+ self.load_network(self.netD, 'D', opt.which_epoch, pretrained_path)
+ if self.gen_features:
+ self.load_network(self.netE, 'E', opt.which_epoch, pretrained_path)
+
+ # set loss functions and optimizers
+ if self.isTrain:
+ if opt.pool_size > 0 and (len(self.gpu_ids)) > 1:
+ raise NotImplementedError("Fake Pool Not Implemented for MultiGPU")
+ self.fake_pool = ImagePool(opt.pool_size)
+ self.old_lr = opt.lr
+
+ # define loss functions
+ self.loss_filter = self.init_loss_filter(not opt.no_ganFeat_loss, not opt.no_vgg_loss)
+
+ self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor)
+ self.criterionFeat = torch.nn.L1Loss()
+ if not opt.no_vgg_loss:
+ self.criterionVGG = networks.VGGLoss(self.gpu_ids)
+
+
+ # Names so we can breakout loss
+ self.loss_names = self.loss_filter('G_GAN','G_GAN_Feat','G_VGG','D_real', 'D_fake')
+
+ # initialize optimizers
+ # optimizer G
+ if opt.niter_fix_global > 0:
+ import sys
+ if sys.version_info >= (3,0):
+ finetune_list = set()
+ else:
+ from sets import Set
+ finetune_list = Set()
+
+ params_dict = dict(self.netG.named_parameters())
+ params = []
+ for key, value in params_dict.items():
+ if key.startswith('model' + str(opt.n_local_enhancers)):
+ params += [value]
+ finetune_list.add(key.split('.')[0])
+ print('------------- Only training the local enhancer network (for %d epochs) ------------' % opt.niter_fix_global)
+ print('The layers that are finetuned are ', sorted(finetune_list))
+ else:
+ params = list(self.netG.parameters())
+ if self.gen_features:
+ params += list(self.netE.parameters())
+ self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+
+ # optimizer D
+ params = list(self.netD.parameters())
+ self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+
+ def encode_input(self, label_map, inst_map=None, real_image=None, feat_map=None, infer=False):
+ if self.opt.label_nc == 0:
+ input_label = label_map.data.cuda()
+ else:
+ # create one-hot vector for label map
+ size = label_map.size()
+ oneHot_size = (size[0], self.opt.label_nc, size[2], size[3])
+ input_label = torch.cuda.FloatTensor(torch.Size(oneHot_size)).zero_()
+ input_label = input_label.scatter_(1, label_map.data.long().cuda(), 1.0)
+ if self.opt.data_type == 16:
+ input_label = input_label.half()
+
+ # get edges from instance map
+ if not self.opt.no_instance:
+ inst_map = inst_map.data.cuda()
+ edge_map = self.get_edges(inst_map)
+ input_label = torch.cat((input_label, edge_map), dim=1)
+ input_label = Variable(input_label, volatile=infer)
+
+ # real images for training
+ if real_image is not None:
+ real_image = Variable(real_image.data.cuda())
+
+ # instance map for feature encoding
+ if self.use_features:
+ # get precomputed feature maps
+ if self.opt.load_features:
+ feat_map = Variable(feat_map.data.cuda())
+ if self.opt.label_feat:
+ inst_map = label_map.cuda()
+
+ return input_label, inst_map, real_image, feat_map
+
+ def discriminate(self, input_label, test_image, use_pool=False):
+ input_concat = torch.cat((input_label, test_image.detach()), dim=1)
+ if use_pool:
+ fake_query = self.fake_pool.query(input_concat)
+ return self.netD.forward(fake_query)
+ else:
+ return self.netD.forward(input_concat)
+
+ def forward(self, label, inst, image, feat, infer=False):
+ # Encode Inputs
+ input_label, inst_map, real_image, feat_map = self.encode_input(label, inst, image, feat)
+
+ # Fake Generation
+ if self.use_features:
+ if not self.opt.load_features:
+ feat_map = self.netE.forward(real_image, inst_map)
+ input_concat = torch.cat((input_label, feat_map), dim=1)
+ else:
+ input_concat = input_label
+ fake_image = self.netG.forward(input_concat)
+
+ # Fake Detection and Loss
+ pred_fake_pool = self.discriminate(input_label, fake_image, use_pool=True)
+ loss_D_fake = self.criterionGAN(pred_fake_pool, False)
+
+ # Real Detection and Loss
+ pred_real = self.discriminate(input_label, real_image)
+ loss_D_real = self.criterionGAN(pred_real, True)
+
+ # GAN loss (Fake Passability Loss)
+ pred_fake = self.netD.forward(torch.cat((input_label, fake_image), dim=1))
+ loss_G_GAN = self.criterionGAN(pred_fake, True)
+
+ # GAN feature matching loss
+ loss_G_GAN_Feat = 0
+ if not self.opt.no_ganFeat_loss:
+ feat_weights = 4.0 / (self.opt.n_layers_D + 1)
+ D_weights = 1.0 / self.opt.num_D
+ for i in range(self.opt.num_D):
+ for j in range(len(pred_fake[i])-1):
+ loss_G_GAN_Feat += D_weights * feat_weights * \
+ self.criterionFeat(pred_fake[i][j], pred_real[i][j].detach()) * self.opt.lambda_feat
+
+ # VGG feature matching loss
+ loss_G_VGG = 0
+ if not self.opt.no_vgg_loss:
+ loss_G_VGG = self.criterionVGG(fake_image, real_image) * self.opt.lambda_feat
+
+ # Only return the fake_B image if necessary to save BW
+ return [ self.loss_filter( loss_G_GAN, loss_G_GAN_Feat, loss_G_VGG, loss_D_real, loss_D_fake ), None if not infer else fake_image ]
+
+ def inference(self, label, inst, image=None):
+ # Encode Inputs
+ image = Variable(image) if image is not None else None
+ input_label, inst_map, real_image, _ = self.encode_input(Variable(label), Variable(inst), image, infer=True)
+
+ # Fake Generation
+ if self.use_features:
+ if self.opt.use_encoded_image:
+ # encode the real image to get feature map
+ feat_map = self.netE.forward(real_image, inst_map)
+ else:
+ # sample clusters from precomputed features
+ feat_map = self.sample_features(inst_map)
+ input_concat = torch.cat((input_label, feat_map), dim=1)
+ else:
+ input_concat = input_label
+
+ if torch.__version__.startswith('0.4'):
+ with torch.no_grad():
+ fake_image = self.netG.forward(input_concat)
+ else:
+ fake_image = self.netG.forward(input_concat)
+ return fake_image
+
+ def sample_features(self, inst):
+ # read precomputed feature clusters
+ cluster_path = os.path.join(self.opt.checkpoints_dir, self.opt.name, self.opt.cluster_path)
+ features_clustered = np.load(cluster_path, encoding='latin1').item()
+
+ # randomly sample from the feature clusters
+ inst_np = inst.cpu().numpy().astype(int)
+ feat_map = self.Tensor(inst.size()[0], self.opt.feat_num, inst.size()[2], inst.size()[3])
+ for i in np.unique(inst_np):
+ label = i if i < 1000 else i//1000
+ if label in features_clustered:
+ feat = features_clustered[label]
+ cluster_idx = np.random.randint(0, feat.shape[0])
+
+ idx = (inst == int(i)).nonzero()
+ for k in range(self.opt.feat_num):
+ feat_map[idx[:,0], idx[:,1] + k, idx[:,2], idx[:,3]] = feat[cluster_idx, k]
+ if self.opt.data_type==16:
+ feat_map = feat_map.half()
+ return feat_map
+
+ def encode_features(self, image, inst):
+ image = Variable(image.cuda(), volatile=True)
+ feat_num = self.opt.feat_num
+ h, w = inst.size()[2], inst.size()[3]
+ block_num = 32
+ feat_map = self.netE.forward(image, inst.cuda())
+ inst_np = inst.cpu().numpy().astype(int)
+ feature = {}
+ for i in range(self.opt.label_nc):
+ feature[i] = np.zeros((0, feat_num+1))
+ for i in np.unique(inst_np):
+ label = i if i < 1000 else i//1000
+ idx = (inst == int(i)).nonzero()
+ num = idx.size()[0]
+ idx = idx[num//2,:]
+ val = np.zeros((1, feat_num+1))
+ for k in range(feat_num):
+ val[0, k] = feat_map[idx[0], idx[1] + k, idx[2], idx[3]].data[0]
+ val[0, feat_num] = float(num) / (h * w // block_num)
+ feature[label] = np.append(feature[label], val, axis=0)
+ return feature
+
+ def get_edges(self, t):
+ edge = torch.cuda.ByteTensor(t.size()).zero_()
+ edge[:,:,:,1:] = edge[:,:,:,1:] | (t[:,:,:,1:] != t[:,:,:,:-1])
+ edge[:,:,:,:-1] = edge[:,:,:,:-1] | (t[:,:,:,1:] != t[:,:,:,:-1])
+ edge[:,:,1:,:] = edge[:,:,1:,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
+ edge[:,:,:-1,:] = edge[:,:,:-1,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
+ if self.opt.data_type==16:
+ return edge.half()
+ else:
+ return edge.float()
+
+ def save(self, which_epoch):
+ self.save_network(self.netG, 'G', which_epoch, self.gpu_ids)
+ self.save_network(self.netD, 'D', which_epoch, self.gpu_ids)
+ if self.gen_features:
+ self.save_network(self.netE, 'E', which_epoch, self.gpu_ids)
+
+ def update_fixed_params(self):
+ # after fixing the global generator for a number of iterations, also start finetuning it
+ params = list(self.netG.parameters())
+ if self.gen_features:
+ params += list(self.netE.parameters())
+ self.optimizer_G = torch.optim.Adam(params, lr=self.opt.lr, betas=(self.opt.beta1, 0.999))
+ if self.opt.verbose:
+ print('------------ Now also finetuning global generator -----------')
+
+ def update_learning_rate(self):
+ lrd = self.opt.lr / self.opt.niter_decay
+ lr = self.old_lr - lrd
+ for param_group in self.optimizer_D.param_groups:
+ param_group['lr'] = lr
+ for param_group in self.optimizer_G.param_groups:
+ param_group['lr'] = lr
+ if self.opt.verbose:
+ print('update learning rate: %f -> %f' % (self.old_lr, lr))
+ self.old_lr = lr
+
+class InferenceModel(Pix2PixHDModel):
+ def forward(self, inp):
+ label, inst = inp
+ return self.inference(label, inst)
+
+
diff --git a/SimSwap/models/projected_model.py b/SimSwap/models/projected_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..15d70f9a6ba5481fc236988d13229f8e4eee31e1
--- /dev/null
+++ b/SimSwap/models/projected_model.py
@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: fs_model_fix_idnorm_donggp_saveoptim copy.py
+# Created Date: Wednesday January 12th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified: Saturday, 13th May 2023 9:56:35 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+
+import torch
+import torch.nn as nn
+
+from .base_model import BaseModel
+from .fs_networks_fix import Generator_Adain_Upsample
+
+from pg_modules.projected_discriminator import ProjectedDiscriminator
+
+def compute_grad2(d_out, x_in):
+ batch_size = x_in.size(0)
+ grad_dout = torch.autograd.grad(
+ outputs=d_out.sum(), inputs=x_in,
+ create_graph=True, retain_graph=True, only_inputs=True
+ )[0]
+ grad_dout2 = grad_dout.pow(2)
+ assert(grad_dout2.size() == x_in.size())
+ reg = grad_dout2.view(batch_size, -1).sum(1)
+ return reg
+
+class fsModel(BaseModel):
+ def name(self):
+ return 'fsModel'
+
+ def initialize(self, opt):
+ BaseModel.initialize(self, opt)
+ # if opt.resize_or_crop != 'none' or not opt.isTrain: # when training at full res this causes OOM
+ self.isTrain = opt.isTrain
+
+ # Generator network
+ self.netG = Generator_Adain_Upsample(input_nc=3, output_nc=3, latent_size=512, n_blocks=9, deep=opt.Gdeep)
+ self.netG.cuda()
+
+ # Id network
+ netArc_checkpoint = opt.Arc_path
+ netArc_checkpoint = torch.load(netArc_checkpoint, map_location=torch.device("cpu"))
+ self.netArc = netArc_checkpoint
+ self.netArc = self.netArc.cuda()
+ self.netArc.eval()
+ self.netArc.requires_grad_(False)
+ if not self.isTrain:
+ pretrained_path = opt.checkpoints_dir
+ self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)
+ return
+ self.netD = ProjectedDiscriminator(diffaug=False, interp224=False, **{})
+ # self.netD.feature_network.requires_grad_(False)
+ self.netD.cuda()
+
+
+ if self.isTrain:
+ # define loss functions
+ self.criterionFeat = nn.L1Loss()
+ self.criterionRec = nn.L1Loss()
+
+
+ # initialize optimizers
+
+ # optimizer G
+ params = list(self.netG.parameters())
+ self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.99),eps=1e-8)
+
+ # optimizer D
+ params = list(self.netD.parameters())
+ self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.99),eps=1e-8)
+
+ # load networks
+ if opt.continue_train:
+ pretrained_path = '' if not self.isTrain else opt.load_pretrain
+ # print (pretrained_path)
+ self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)
+ self.load_network(self.netD, 'D', opt.which_epoch, pretrained_path)
+ self.load_optim(self.optimizer_G, 'G', opt.which_epoch, pretrained_path)
+ self.load_optim(self.optimizer_D, 'D', opt.which_epoch, pretrained_path)
+ torch.cuda.empty_cache()
+
+ def cosin_metric(self, x1, x2):
+ #return np.dot(x1, x2) / (np.linalg.norm(x1) * np.linalg.norm(x2))
+ return torch.sum(x1 * x2, dim=1) / (torch.norm(x1, dim=1) * torch.norm(x2, dim=1))
+
+
+
+ def save(self, which_epoch):
+ self.save_network(self.netG, 'G', which_epoch)
+ self.save_network(self.netD, 'D', which_epoch)
+ self.save_optim(self.optimizer_G, 'G', which_epoch)
+ self.save_optim(self.optimizer_D, 'D', which_epoch)
+ '''if self.gen_features:
+ self.save_network(self.netE, 'E', which_epoch, self.gpu_ids)'''
+
+ def update_fixed_params(self):
+ # after fixing the global generator for a number of iterations, also start finetuning it
+ params = list(self.netG.parameters())
+ if self.gen_features:
+ params += list(self.netE.parameters())
+ self.optimizer_G = torch.optim.Adam(params, lr=self.opt.lr, betas=(self.opt.beta1, 0.999))
+ if self.opt.verbose:
+ print('------------ Now also finetuning global generator -----------')
+
+ def update_learning_rate(self):
+ lrd = self.opt.lr / self.opt.niter_decay
+ lr = self.old_lr - lrd
+ for param_group in self.optimizer_D.param_groups:
+ param_group['lr'] = lr
+ for param_group in self.optimizer_G.param_groups:
+ param_group['lr'] = lr
+ if self.opt.verbose:
+ print('update learning rate: %f -> %f' % (self.old_lr, lr))
+ self.old_lr = lr
+
+
diff --git a/SimSwap/models/projectionhead.py b/SimSwap/models/projectionhead.py
new file mode 100644
index 0000000000000000000000000000000000000000..a0d18f1f25f7dcba73172d3d9f62de42b194987d
--- /dev/null
+++ b/SimSwap/models/projectionhead.py
@@ -0,0 +1,14 @@
+import torch.nn as nn
+
+class ProjectionHead(nn.Module):
+ def __init__(self, proj_dim=256):
+ super(ProjectionHead, self).__init__()
+
+ self.proj = nn.Sequential(
+ nn.Linear(proj_dim, proj_dim),
+ nn.ReLU(),
+ nn.Linear(proj_dim, proj_dim),
+ )
+
+ def forward(self, x):
+ return self.proj(x)
\ No newline at end of file
diff --git a/SimSwap/models/ui_model.py b/SimSwap/models/ui_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..c5b34335493c51fd110c063120cbb961ec9f2014
--- /dev/null
+++ b/SimSwap/models/ui_model.py
@@ -0,0 +1,347 @@
+import torch
+from torch.autograd import Variable
+from collections import OrderedDict
+import numpy as np
+import os
+from PIL import Image
+import util.util as util
+from .base_model import BaseModel
+from . import networks
+
+class UIModel(BaseModel):
+ def name(self):
+ return 'UIModel'
+
+ def initialize(self, opt):
+ assert(not opt.isTrain)
+ BaseModel.initialize(self, opt)
+ self.use_features = opt.instance_feat or opt.label_feat
+
+ netG_input_nc = opt.label_nc
+ if not opt.no_instance:
+ netG_input_nc += 1
+ if self.use_features:
+ netG_input_nc += opt.feat_num
+
+ self.netG = networks.define_G(netG_input_nc, opt.output_nc, opt.ngf, opt.netG,
+ opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers,
+ opt.n_blocks_local, opt.norm, gpu_ids=self.gpu_ids)
+ self.load_network(self.netG, 'G', opt.which_epoch)
+
+ print('---------- Networks initialized -------------')
+
+ def toTensor(self, img, normalize=False):
+ tensor = torch.from_numpy(np.array(img, np.int32, copy=False))
+ tensor = tensor.view(1, img.size[1], img.size[0], len(img.mode))
+ tensor = tensor.transpose(1, 2).transpose(1, 3).contiguous()
+ if normalize:
+ return (tensor.float()/255.0 - 0.5) / 0.5
+ return tensor.float()
+
+ def load_image(self, label_path, inst_path, feat_path):
+ opt = self.opt
+ # read label map
+ label_img = Image.open(label_path)
+ if label_path.find('face') != -1:
+ label_img = label_img.convert('L')
+ ow, oh = label_img.size
+ w = opt.loadSize
+ h = int(w * oh / ow)
+ label_img = label_img.resize((w, h), Image.NEAREST)
+ label_map = self.toTensor(label_img)
+
+ # onehot vector input for label map
+ self.label_map = label_map.cuda()
+ oneHot_size = (1, opt.label_nc, h, w)
+ input_label = self.Tensor(torch.Size(oneHot_size)).zero_()
+ self.input_label = input_label.scatter_(1, label_map.long().cuda(), 1.0)
+
+ # read instance map
+ if not opt.no_instance:
+ inst_img = Image.open(inst_path)
+ inst_img = inst_img.resize((w, h), Image.NEAREST)
+ self.inst_map = self.toTensor(inst_img).cuda()
+ self.edge_map = self.get_edges(self.inst_map)
+ self.net_input = Variable(torch.cat((self.input_label, self.edge_map), dim=1), volatile=True)
+ else:
+ self.net_input = Variable(self.input_label, volatile=True)
+
+ self.features_clustered = np.load(feat_path).item()
+ self.object_map = self.inst_map if opt.instance_feat else self.label_map
+
+ object_np = self.object_map.cpu().numpy().astype(int)
+ self.feat_map = self.Tensor(1, opt.feat_num, h, w).zero_()
+ self.cluster_indices = np.zeros(self.opt.label_nc, np.uint8)
+ for i in np.unique(object_np):
+ label = i if i < 1000 else i//1000
+ if label in self.features_clustered:
+ feat = self.features_clustered[label]
+ np.random.seed(i+1)
+ cluster_idx = np.random.randint(0, feat.shape[0])
+ self.cluster_indices[label] = cluster_idx
+ idx = (self.object_map == i).nonzero()
+ self.set_features(idx, feat, cluster_idx)
+
+ self.net_input_original = self.net_input.clone()
+ self.label_map_original = self.label_map.clone()
+ self.feat_map_original = self.feat_map.clone()
+ if not opt.no_instance:
+ self.inst_map_original = self.inst_map.clone()
+
+ def reset(self):
+ self.net_input = self.net_input_prev = self.net_input_original.clone()
+ self.label_map = self.label_map_prev = self.label_map_original.clone()
+ self.feat_map = self.feat_map_prev = self.feat_map_original.clone()
+ if not self.opt.no_instance:
+ self.inst_map = self.inst_map_prev = self.inst_map_original.clone()
+ self.object_map = self.inst_map if self.opt.instance_feat else self.label_map
+
+ def undo(self):
+ self.net_input = self.net_input_prev
+ self.label_map = self.label_map_prev
+ self.feat_map = self.feat_map_prev
+ if not self.opt.no_instance:
+ self.inst_map = self.inst_map_prev
+ self.object_map = self.inst_map if self.opt.instance_feat else self.label_map
+
+ # get boundary map from instance map
+ def get_edges(self, t):
+ edge = torch.cuda.ByteTensor(t.size()).zero_()
+ edge[:,:,:,1:] = edge[:,:,:,1:] | (t[:,:,:,1:] != t[:,:,:,:-1])
+ edge[:,:,:,:-1] = edge[:,:,:,:-1] | (t[:,:,:,1:] != t[:,:,:,:-1])
+ edge[:,:,1:,:] = edge[:,:,1:,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
+ edge[:,:,:-1,:] = edge[:,:,:-1,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
+ return edge.float()
+
+ # change the label at the source position to the label at the target position
+ def change_labels(self, click_src, click_tgt):
+ y_src, x_src = click_src[0], click_src[1]
+ y_tgt, x_tgt = click_tgt[0], click_tgt[1]
+ label_src = int(self.label_map[0, 0, y_src, x_src])
+ inst_src = self.inst_map[0, 0, y_src, x_src]
+ label_tgt = int(self.label_map[0, 0, y_tgt, x_tgt])
+ inst_tgt = self.inst_map[0, 0, y_tgt, x_tgt]
+
+ idx_src = (self.inst_map == inst_src).nonzero()
+ # need to change 3 things: label map, instance map, and feature map
+ if idx_src.shape:
+ # backup current maps
+ self.backup_current_state()
+
+ # change both the label map and the network input
+ self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
+ self.net_input[idx_src[:,0], idx_src[:,1] + label_src, idx_src[:,2], idx_src[:,3]] = 0
+ self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1
+
+ # update the instance map (and the network input)
+ if inst_tgt > 1000:
+ # if different instances have different ids, give the new object a new id
+ tgt_indices = (self.inst_map > label_tgt * 1000) & (self.inst_map < (label_tgt+1) * 1000)
+ inst_tgt = self.inst_map[tgt_indices].max() + 1
+ self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = inst_tgt
+ self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)
+
+ # also copy the source features to the target position
+ idx_tgt = (self.inst_map == inst_tgt).nonzero()
+ if idx_tgt.shape:
+ self.copy_features(idx_src, idx_tgt[0,:])
+
+ self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
+
+ # add strokes of target label in the image
+ def add_strokes(self, click_src, label_tgt, bw, save):
+ # get the region of the new strokes (bw is the brush width)
+ size = self.net_input.size()
+ h, w = size[2], size[3]
+ idx_src = torch.LongTensor(bw**2, 4).fill_(0)
+ for i in range(bw):
+ idx_src[i*bw:(i+1)*bw, 2] = min(h-1, max(0, click_src[0]-bw//2 + i))
+ for j in range(bw):
+ idx_src[i*bw+j, 3] = min(w-1, max(0, click_src[1]-bw//2 + j))
+ idx_src = idx_src.cuda()
+
+ # again, need to update 3 things
+ if idx_src.shape:
+ # backup current maps
+ if save:
+ self.backup_current_state()
+
+ # update the label map (and the network input) in the stroke region
+ self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
+ for k in range(self.opt.label_nc):
+ self.net_input[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = 0
+ self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1
+
+ # update the instance map (and the network input)
+ self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
+ self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)
+
+ # also update the features if available
+ if self.opt.instance_feat:
+ feat = self.features_clustered[label_tgt]
+ #np.random.seed(label_tgt+1)
+ #cluster_idx = np.random.randint(0, feat.shape[0])
+ cluster_idx = self.cluster_indices[label_tgt]
+ self.set_features(idx_src, feat, cluster_idx)
+
+ self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
+
+ # add an object to the clicked position with selected style
+ def add_objects(self, click_src, label_tgt, mask, style_id=0):
+ y, x = click_src[0], click_src[1]
+ mask = np.transpose(mask, (2, 0, 1))[np.newaxis,...]
+ idx_src = torch.from_numpy(mask).cuda().nonzero()
+ idx_src[:,2] += y
+ idx_src[:,3] += x
+
+ # backup current maps
+ self.backup_current_state()
+
+ # update label map
+ self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
+ for k in range(self.opt.label_nc):
+ self.net_input[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = 0
+ self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1
+
+ # update instance map
+ self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
+ self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)
+
+ # update feature map
+ self.set_features(idx_src, self.feat, style_id)
+
+ self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
+
+ def single_forward(self, net_input, feat_map):
+ net_input = torch.cat((net_input, feat_map), dim=1)
+ fake_image = self.netG.forward(net_input)
+
+ if fake_image.size()[0] == 1:
+ return fake_image.data[0]
+ return fake_image.data
+
+
+ # generate all outputs for different styles
+ def style_forward(self, click_pt, style_id=-1):
+ if click_pt is None:
+ self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
+ self.crop = None
+ self.mask = None
+ else:
+ instToChange = int(self.object_map[0, 0, click_pt[0], click_pt[1]])
+ self.instToChange = instToChange
+ label = instToChange if instToChange < 1000 else instToChange//1000
+ self.feat = self.features_clustered[label]
+ self.fake_image = []
+ self.mask = self.object_map == instToChange
+ idx = self.mask.nonzero()
+ self.get_crop_region(idx)
+ if idx.size():
+ if style_id == -1:
+ (min_y, min_x, max_y, max_x) = self.crop
+ ### original
+ for cluster_idx in range(self.opt.multiple_output):
+ self.set_features(idx, self.feat, cluster_idx)
+ fake_image = self.single_forward(self.net_input, self.feat_map)
+ fake_image = util.tensor2im(fake_image[:,min_y:max_y,min_x:max_x])
+ self.fake_image.append(fake_image)
+ """### To speed up previewing different style results, either crop or downsample the label maps
+ if instToChange > 1000:
+ (min_y, min_x, max_y, max_x) = self.crop
+ ### crop
+ _, _, h, w = self.net_input.size()
+ offset = 512
+ y_start, x_start = max(0, min_y-offset), max(0, min_x-offset)
+ y_end, x_end = min(h, (max_y + offset)), min(w, (max_x + offset))
+ y_region = slice(y_start, y_start+(y_end-y_start)//16*16)
+ x_region = slice(x_start, x_start+(x_end-x_start)//16*16)
+ net_input = self.net_input[:,:,y_region,x_region]
+ for cluster_idx in range(self.opt.multiple_output):
+ self.set_features(idx, self.feat, cluster_idx)
+ fake_image = self.single_forward(net_input, self.feat_map[:,:,y_region,x_region])
+ fake_image = util.tensor2im(fake_image[:,min_y-y_start:max_y-y_start,min_x-x_start:max_x-x_start])
+ self.fake_image.append(fake_image)
+ else:
+ ### downsample
+ (min_y, min_x, max_y, max_x) = [crop//2 for crop in self.crop]
+ net_input = self.net_input[:,:,::2,::2]
+ size = net_input.size()
+ net_input_batch = net_input.expand(self.opt.multiple_output, size[1], size[2], size[3])
+ for cluster_idx in range(self.opt.multiple_output):
+ self.set_features(idx, self.feat, cluster_idx)
+ feat_map = self.feat_map[:,:,::2,::2]
+ if cluster_idx == 0:
+ feat_map_batch = feat_map
+ else:
+ feat_map_batch = torch.cat((feat_map_batch, feat_map), dim=0)
+ fake_image_batch = self.single_forward(net_input_batch, feat_map_batch)
+ for i in range(self.opt.multiple_output):
+ self.fake_image.append(util.tensor2im(fake_image_batch[i,:,min_y:max_y,min_x:max_x]))"""
+
+ else:
+ self.set_features(idx, self.feat, style_id)
+ self.cluster_indices[label] = style_id
+ self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
+
+ def backup_current_state(self):
+ self.net_input_prev = self.net_input.clone()
+ self.label_map_prev = self.label_map.clone()
+ self.inst_map_prev = self.inst_map.clone()
+ self.feat_map_prev = self.feat_map.clone()
+
+ # crop the ROI and get the mask of the object
+ def get_crop_region(self, idx):
+ size = self.net_input.size()
+ h, w = size[2], size[3]
+ min_y, min_x = idx[:,2].min(), idx[:,3].min()
+ max_y, max_x = idx[:,2].max(), idx[:,3].max()
+ crop_min = 128
+ if max_y - min_y < crop_min:
+ min_y = max(0, (max_y + min_y) // 2 - crop_min // 2)
+ max_y = min(h-1, min_y + crop_min)
+ if max_x - min_x < crop_min:
+ min_x = max(0, (max_x + min_x) // 2 - crop_min // 2)
+ max_x = min(w-1, min_x + crop_min)
+ self.crop = (min_y, min_x, max_y, max_x)
+ self.mask = self.mask[:,:, min_y:max_y, min_x:max_x]
+
+ # update the feature map once a new object is added or the label is changed
+ def update_features(self, cluster_idx, mask=None, click_pt=None):
+ self.feat_map_prev = self.feat_map.clone()
+ # adding a new object
+ if mask is not None:
+ y, x = click_pt[0], click_pt[1]
+ mask = np.transpose(mask, (2,0,1))[np.newaxis,...]
+ idx = torch.from_numpy(mask).cuda().nonzero()
+ idx[:,2] += y
+ idx[:,3] += x
+ # changing the label of an existing object
+ else:
+ idx = (self.object_map == self.instToChange).nonzero()
+
+ # update feature map
+ self.set_features(idx, self.feat, cluster_idx)
+
+ # set the class features to the target feature
+ def set_features(self, idx, feat, cluster_idx):
+ for k in range(self.opt.feat_num):
+ self.feat_map[idx[:,0], idx[:,1] + k, idx[:,2], idx[:,3]] = feat[cluster_idx, k]
+
+ # copy the features at the target position to the source position
+ def copy_features(self, idx_src, idx_tgt):
+ for k in range(self.opt.feat_num):
+ val = self.feat_map[idx_tgt[0], idx_tgt[1] + k, idx_tgt[2], idx_tgt[3]]
+ self.feat_map[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = val
+
+ def get_current_visuals(self, getLabel=False):
+ mask = self.mask
+ if self.mask is not None:
+ mask = np.transpose(self.mask[0].cpu().float().numpy(), (1,2,0)).astype(np.uint8)
+
+ dict_list = [('fake_image', self.fake_image), ('mask', mask)]
+
+ if getLabel: # only output label map if needed to save bandwidth
+ label = util.tensor2label(self.net_input.data[0], self.opt.label_nc)
+ dict_list += [('label', label)]
+
+ return OrderedDict(dict_list)
\ No newline at end of file
diff --git a/SimSwap/options/__pycache__/base_options.cpython-310.pyc b/SimSwap/options/__pycache__/base_options.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9cf4d66db7c2c19525aaac36def7523c85a6b41f
Binary files /dev/null and b/SimSwap/options/__pycache__/base_options.cpython-310.pyc differ
diff --git a/SimSwap/options/__pycache__/test_options.cpython-310.pyc b/SimSwap/options/__pycache__/test_options.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..131eef2bb764b98b8218ead63b9b6260def77101
Binary files /dev/null and b/SimSwap/options/__pycache__/test_options.cpython-310.pyc differ
diff --git a/SimSwap/options/base_options.py b/SimSwap/options/base_options.py
new file mode 100644
index 0000000000000000000000000000000000000000..14f817b10d1ce3711c8020f5ab53a2e862276090
--- /dev/null
+++ b/SimSwap/options/base_options.py
@@ -0,0 +1,104 @@
+import argparse
+import os
+from util import util
+import torch
+
+class BaseOptions():
+ def __init__(self):
+ self.parser = argparse.ArgumentParser()
+ self.initialized = False
+
+ def initialize(self):
+ # experiment specifics
+ self.parser.add_argument('--name', type=str, default='people', help='name of the experiment. It decides where to store samples and models')
+ self.parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU')
+ self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
+ # self.parser.add_argument('--model', type=str, default='pix2pixHD', help='which model to use')
+ self.parser.add_argument('--norm', type=str, default='batch', help='instance normalization or batch normalization')
+ self.parser.add_argument('--use_dropout', action='store_true', help='use dropout for the generator')
+ self.parser.add_argument('--data_type', default=32, type=int, choices=[8, 16, 32], help="Supported data type i.e. 8, 16, 32 bit")
+ self.parser.add_argument('--verbose', action='store_true', default=False, help='toggles verbose')
+ self.parser.add_argument('--fp16', action='store_true', default=False, help='train with AMP')
+ self.parser.add_argument('--local_rank', type=int, default=0, help='local rank for distributed training')
+ self.parser.add_argument('--isTrain', type=bool, default=True, help='local rank for distributed training')
+
+ # input/output sizes
+ self.parser.add_argument('--batchSize', type=int, default=8, help='input batch size')
+ self.parser.add_argument('--loadSize', type=int, default=1024, help='scale images to this size')
+ self.parser.add_argument('--fineSize', type=int, default=512, help='then crop to this size')
+ self.parser.add_argument('--label_nc', type=int, default=0, help='# of input label channels')
+ self.parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels')
+ self.parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels')
+
+ # for setting inputs
+ self.parser.add_argument('--dataroot', type=str, default='./datasets/cityscapes/')
+ self.parser.add_argument('--resize_or_crop', type=str, default='scale_width', help='scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop]')
+ self.parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
+ self.parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data argumentation')
+ self.parser.add_argument('--nThreads', default=2, type=int, help='# threads for loading data')
+ self.parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
+
+ # for displays
+ self.parser.add_argument('--display_winsize', type=int, default=512, help='display window size')
+ self.parser.add_argument('--tf_log', action='store_true', help='if specified, use tensorboard logging. Requires tensorflow installed')
+
+ # for generator
+ self.parser.add_argument('--netG', type=str, default='global', help='selects model to use for netG')
+ self.parser.add_argument('--latent_size', type=int, default=512, help='latent size of Adain layer')
+ self.parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')
+ self.parser.add_argument('--n_downsample_global', type=int, default=3, help='number of downsampling layers in netG')
+ self.parser.add_argument('--n_blocks_global', type=int, default=6, help='number of residual blocks in the global generator network')
+ self.parser.add_argument('--n_blocks_local', type=int, default=3, help='number of residual blocks in the local enhancer network')
+ self.parser.add_argument('--n_local_enhancers', type=int, default=1, help='number of local enhancers to use')
+ self.parser.add_argument('--niter_fix_global', type=int, default=0, help='number of epochs that we only train the outmost local enhancer')
+
+ # for instance-wise features
+ self.parser.add_argument('--no_instance', action='store_true', help='if specified, do *not* add instance map as input')
+ self.parser.add_argument('--instance_feat', action='store_true', help='if specified, add encoded instance features as input')
+ self.parser.add_argument('--label_feat', action='store_true', help='if specified, add encoded label features as input')
+ self.parser.add_argument('--feat_num', type=int, default=3, help='vector length for encoded features')
+ self.parser.add_argument('--load_features', action='store_true', help='if specified, load precomputed feature maps')
+ self.parser.add_argument('--n_downsample_E', type=int, default=4, help='# of downsampling layers in encoder')
+ self.parser.add_argument('--nef', type=int, default=16, help='# of encoder filters in the first conv layer')
+ self.parser.add_argument('--n_clusters', type=int, default=10, help='number of clusters for features')
+ self.parser.add_argument('--image_size', type=int, default=224, help='number of clusters for features')
+ self.parser.add_argument('--norm_G', type=str, default='spectralspadesyncbatch3x3', help='instance normalization or batch normalization')
+ self.parser.add_argument('--semantic_nc', type=int, default=3, help='number of clusters for features')
+ self.initialized = True
+
+ def parse(self, save=True):
+ if not self.initialized:
+ self.initialize()
+ self.opt = self.parser.parse_args()
+ self.opt.isTrain = self.isTrain # train or test
+
+ str_ids = self.opt.gpu_ids.split(',')
+ self.opt.gpu_ids = []
+ for str_id in str_ids:
+ id = int(str_id)
+ if id >= 0:
+ self.opt.gpu_ids.append(id)
+
+ # set gpu ids
+ if len(self.opt.gpu_ids) > 0:
+ torch.cuda.set_device(self.opt.gpu_ids[0])
+
+ args = vars(self.opt)
+
+ print('------------ Options -------------')
+ for k, v in sorted(args.items()):
+ print('%s: %s' % (str(k), str(v)))
+ print('-------------- End ----------------')
+
+ # save to the disk
+ if self.opt.isTrain:
+ expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name)
+ util.mkdirs(expr_dir)
+ if save and not self.opt.continue_train:
+ file_name = os.path.join(expr_dir, 'opt.txt')
+ with open(file_name, 'wt') as opt_file:
+ opt_file.write('------------ Options -------------\n')
+ for k, v in sorted(args.items()):
+ opt_file.write('%s: %s\n' % (str(k), str(v)))
+ opt_file.write('-------------- End ----------------\n')
+ return self.opt
diff --git a/SimSwap/options/test_options.py b/SimSwap/options/test_options.py
new file mode 100644
index 0000000000000000000000000000000000000000..3562671b4f966585dae18d57145c7950fbabf4e2
--- /dev/null
+++ b/SimSwap/options/test_options.py
@@ -0,0 +1,38 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 17:03:58
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-23 17:08:08
+Description:
+'''
+from .base_options import BaseOptions
+
+class TestOptions(BaseOptions):
+ def initialize(self):
+ BaseOptions.initialize(self)
+ self.parser.add_argument('--ntest', type=int, default=float("inf"), help='# of test examples.')
+ self.parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
+ self.parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')
+ self.parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
+ self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
+ self.parser.add_argument('--how_many', type=int, default=50, help='how many test images to run')
+ self.parser.add_argument('--cluster_path', type=str, default='features_clustered_010.npy', help='the path for clustered results of encoded features')
+ self.parser.add_argument('--use_encoded_image', action='store_true', help='if specified, encode the real image to get the feature map')
+ self.parser.add_argument("--export_onnx", type=str, help="export ONNX model to a given file")
+ self.parser.add_argument("--engine", type=str, help="run serialized TRT engine")
+ self.parser.add_argument("--onnx", type=str, help="run ONNX model via TRT")
+ self.parser.add_argument("--Arc_path", type=str, default='arcface_model/arcface_checkpoint.tar', help="run ONNX model via TRT")
+ self.parser.add_argument("--pic_a_path", type=str, default='G:/swap_data/ID/elon-musk-hero-image.jpeg', help="Person who provides identity information")
+ self.parser.add_argument("--pic_b_path", type=str, default='./demo_file/multi_people.jpg', help="Person who provides information other than their identity")
+ self.parser.add_argument("--pic_specific_path", type=str, default='./crop_224/zrf.jpg', help="The specific person to be swapped")
+ self.parser.add_argument("--multisepcific_dir", type=str, default='./demo_file/multispecific', help="Dir for multi specific")
+ self.parser.add_argument("--video_path", type=str, default='G:/swap_data/video/HSB_Demo_Trim.mp4', help="path for the video to swap")
+ self.parser.add_argument("--temp_path", type=str, default='./temp_results', help="path to save temporarily images")
+ self.parser.add_argument("--output_path", type=str, default='./output/', help="results path")
+ self.parser.add_argument('--id_thres', type=float, default=0.03, help='how many test images to run')
+ self.parser.add_argument('--no_simswaplogo', action='store_true', help='Remove the watermark')
+ self.parser.add_argument('--use_mask', action='store_true', help='Use mask for better result')
+ self.parser.add_argument('--crop_size', type=int, default=512, help='Crop of size of input image')
+
+ self.isTrain = False
\ No newline at end of file
diff --git a/SimSwap/options/train_options.py b/SimSwap/options/train_options.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a34357f17dc364eadcf55c9fa7bda6ac91e67a2
--- /dev/null
+++ b/SimSwap/options/train_options.py
@@ -0,0 +1,39 @@
+from .base_options import BaseOptions
+
+class TrainOptions(BaseOptions):
+ def initialize(self):
+ BaseOptions.initialize(self)
+ # for displays
+ self.parser.add_argument('--display_freq', type=int, default=99, help='frequency of showing training results on screen')
+ self.parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')
+ self.parser.add_argument('--save_latest_freq', type=int, default=10000, help='frequency of saving the latest results')
+ self.parser.add_argument('--save_epoch_freq', type=int, default=10000, help='frequency of saving checkpoints at the end of epochs')
+ self.parser.add_argument('--no_html', action='store_true', help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/')
+ self.parser.add_argument('--debug', action='store_true', help='only do one epoch and displays at each iteration')
+
+ # for training
+ self.parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
+ self.parser.add_argument('--load_pretrain', type=str, default='', help='load the pretrained model from the specified location')
+ self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
+ self.parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
+ self.parser.add_argument('--niter', type=int, default=10000, help='# of iter at starting learning rate')
+ self.parser.add_argument('--niter_decay', type=int, default=10000, help='# of iter to linearly decay learning rate to zero')
+ self.parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')
+ self.parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')
+
+ # for discriminators
+ self.parser.add_argument('--num_D', type=int, default=2, help='number of discriminators to use')
+ self.parser.add_argument('--n_layers_D', type=int, default=4, help='only used if which_model_netD==n_layers')
+ self.parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer')
+ self.parser.add_argument('--lambda_feat', type=float, default=10.0, help='weight for feature matching loss')
+ self.parser.add_argument('--lambda_id', type=float, default=20.0, help='weight for id loss')
+ self.parser.add_argument('--lambda_rec', type=float, default=10.0, help='weight for reconstruction loss')
+ self.parser.add_argument('--lambda_GP', type=float, default=10.0, help='weight for gradient penalty loss')
+ self.parser.add_argument('--no_ganFeat_loss', action='store_true', help='if specified, do *not* use discriminator feature matching loss')
+ self.parser.add_argument('--no_vgg_loss', action='store_true', help='if specified, do *not* use VGG feature matching loss')
+ self.parser.add_argument('--gan_mode', type=str, default='hinge', help='(ls|original|hinge)')
+ self.parser.add_argument('--pool_size', type=int, default=0, help='the size of image buffer that stores previously generated images')
+ self.parser.add_argument('--times_G', type=int, default=1,
+ help='time of training generator before traning discriminator')
+
+ self.isTrain = True
diff --git a/SimSwap/parsing_model/__pycache__/model.cpython-310.pyc b/SimSwap/parsing_model/__pycache__/model.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c8ef349d0e0d40d35a2eb411e603b6700d6fe9e2
Binary files /dev/null and b/SimSwap/parsing_model/__pycache__/model.cpython-310.pyc differ
diff --git a/SimSwap/parsing_model/__pycache__/resnet.cpython-310.pyc b/SimSwap/parsing_model/__pycache__/resnet.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4888feb33e2570e185d79100387016110ffdd307
Binary files /dev/null and b/SimSwap/parsing_model/__pycache__/resnet.cpython-310.pyc differ
diff --git a/SimSwap/parsing_model/checkpoint/79999_iter.pth b/SimSwap/parsing_model/checkpoint/79999_iter.pth
new file mode 100644
index 0000000000000000000000000000000000000000..ca57f3257ca7715bc340d065764bc249d985c287
--- /dev/null
+++ b/SimSwap/parsing_model/checkpoint/79999_iter.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:468e13ca13a9b43cc0881a9f99083a430e9c0a38abd935431d1c28ee94b26567
+size 53289463
diff --git a/SimSwap/parsing_model/model.py b/SimSwap/parsing_model/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..d434836adf38f61a6e8a05b4ffb29421342580df
--- /dev/null
+++ b/SimSwap/parsing_model/model.py
@@ -0,0 +1,283 @@
+#!/usr/bin/python
+# -*- encoding: utf-8 -*-
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+from parsing_model.resnet import Resnet18
+# from modules.bn import InPlaceABNSync as BatchNorm2d
+
+
+class ConvBNReLU(nn.Module):
+ def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs):
+ super(ConvBNReLU, self).__init__()
+ self.conv = nn.Conv2d(in_chan,
+ out_chan,
+ kernel_size = ks,
+ stride = stride,
+ padding = padding,
+ bias = False)
+ self.bn = nn.BatchNorm2d(out_chan)
+ self.init_weight()
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = F.relu(self.bn(x))
+ return x
+
+ def init_weight(self):
+ for ly in self.children():
+ if isinstance(ly, nn.Conv2d):
+ nn.init.kaiming_normal_(ly.weight, a=1)
+ if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+class BiSeNetOutput(nn.Module):
+ def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs):
+ super(BiSeNetOutput, self).__init__()
+ self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)
+ self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False)
+ self.init_weight()
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.conv_out(x)
+ return x
+
+ def init_weight(self):
+ for ly in self.children():
+ if isinstance(ly, nn.Conv2d):
+ nn.init.kaiming_normal_(ly.weight, a=1)
+ if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+ def get_params(self):
+ wd_params, nowd_params = [], []
+ for name, module in self.named_modules():
+ if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+ wd_params.append(module.weight)
+ if not module.bias is None:
+ nowd_params.append(module.bias)
+ elif isinstance(module, nn.BatchNorm2d):
+ nowd_params += list(module.parameters())
+ return wd_params, nowd_params
+
+
+class AttentionRefinementModule(nn.Module):
+ def __init__(self, in_chan, out_chan, *args, **kwargs):
+ super(AttentionRefinementModule, self).__init__()
+ self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
+ self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False)
+ self.bn_atten = nn.BatchNorm2d(out_chan)
+ self.sigmoid_atten = nn.Sigmoid()
+ self.init_weight()
+
+ def forward(self, x):
+ feat = self.conv(x)
+ atten = F.avg_pool2d(feat, feat.size()[2:])
+ atten = self.conv_atten(atten)
+ atten = self.bn_atten(atten)
+ atten = self.sigmoid_atten(atten)
+ out = torch.mul(feat, atten)
+ return out
+
+ def init_weight(self):
+ for ly in self.children():
+ if isinstance(ly, nn.Conv2d):
+ nn.init.kaiming_normal_(ly.weight, a=1)
+ if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+
+class ContextPath(nn.Module):
+ def __init__(self, *args, **kwargs):
+ super(ContextPath, self).__init__()
+ self.resnet = Resnet18()
+ self.arm16 = AttentionRefinementModule(256, 128)
+ self.arm32 = AttentionRefinementModule(512, 128)
+ self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+ self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+ self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)
+
+ self.init_weight()
+
+ def forward(self, x):
+ H0, W0 = x.size()[2:]
+ feat8, feat16, feat32 = self.resnet(x)
+ H8, W8 = feat8.size()[2:]
+ H16, W16 = feat16.size()[2:]
+ H32, W32 = feat32.size()[2:]
+
+ avg = F.avg_pool2d(feat32, feat32.size()[2:])
+ avg = self.conv_avg(avg)
+ avg_up = F.interpolate(avg, (H32, W32), mode='nearest')
+
+ feat32_arm = self.arm32(feat32)
+ feat32_sum = feat32_arm + avg_up
+ feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest')
+ feat32_up = self.conv_head32(feat32_up)
+
+ feat16_arm = self.arm16(feat16)
+ feat16_sum = feat16_arm + feat32_up
+ feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest')
+ feat16_up = self.conv_head16(feat16_up)
+
+ return feat8, feat16_up, feat32_up # x8, x8, x16
+
+ def init_weight(self):
+ for ly in self.children():
+ if isinstance(ly, nn.Conv2d):
+ nn.init.kaiming_normal_(ly.weight, a=1)
+ if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+ def get_params(self):
+ wd_params, nowd_params = [], []
+ for name, module in self.named_modules():
+ if isinstance(module, (nn.Linear, nn.Conv2d)):
+ wd_params.append(module.weight)
+ if not module.bias is None:
+ nowd_params.append(module.bias)
+ elif isinstance(module, nn.BatchNorm2d):
+ nowd_params += list(module.parameters())
+ return wd_params, nowd_params
+
+
+### This is not used, since I replace this with the resnet feature with the same size
+class SpatialPath(nn.Module):
+ def __init__(self, *args, **kwargs):
+ super(SpatialPath, self).__init__()
+ self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3)
+ self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
+ self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
+ self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0)
+ self.init_weight()
+
+ def forward(self, x):
+ feat = self.conv1(x)
+ feat = self.conv2(feat)
+ feat = self.conv3(feat)
+ feat = self.conv_out(feat)
+ return feat
+
+ def init_weight(self):
+ for ly in self.children():
+ if isinstance(ly, nn.Conv2d):
+ nn.init.kaiming_normal_(ly.weight, a=1)
+ if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+ def get_params(self):
+ wd_params, nowd_params = [], []
+ for name, module in self.named_modules():
+ if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+ wd_params.append(module.weight)
+ if not module.bias is None:
+ nowd_params.append(module.bias)
+ elif isinstance(module, nn.BatchNorm2d):
+ nowd_params += list(module.parameters())
+ return wd_params, nowd_params
+
+
+class FeatureFusionModule(nn.Module):
+ def __init__(self, in_chan, out_chan, *args, **kwargs):
+ super(FeatureFusionModule, self).__init__()
+ self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)
+ self.conv1 = nn.Conv2d(out_chan,
+ out_chan//4,
+ kernel_size = 1,
+ stride = 1,
+ padding = 0,
+ bias = False)
+ self.conv2 = nn.Conv2d(out_chan//4,
+ out_chan,
+ kernel_size = 1,
+ stride = 1,
+ padding = 0,
+ bias = False)
+ self.relu = nn.ReLU(inplace=True)
+ self.sigmoid = nn.Sigmoid()
+ self.init_weight()
+
+ def forward(self, fsp, fcp):
+ fcat = torch.cat([fsp, fcp], dim=1)
+ feat = self.convblk(fcat)
+ atten = F.avg_pool2d(feat, feat.size()[2:])
+ atten = self.conv1(atten)
+ atten = self.relu(atten)
+ atten = self.conv2(atten)
+ atten = self.sigmoid(atten)
+ feat_atten = torch.mul(feat, atten)
+ feat_out = feat_atten + feat
+ return feat_out
+
+ def init_weight(self):
+ for ly in self.children():
+ if isinstance(ly, nn.Conv2d):
+ nn.init.kaiming_normal_(ly.weight, a=1)
+ if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+ def get_params(self):
+ wd_params, nowd_params = [], []
+ for name, module in self.named_modules():
+ if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+ wd_params.append(module.weight)
+ if not module.bias is None:
+ nowd_params.append(module.bias)
+ elif isinstance(module, nn.BatchNorm2d):
+ nowd_params += list(module.parameters())
+ return wd_params, nowd_params
+
+
+class BiSeNet(nn.Module):
+ def __init__(self, n_classes, *args, **kwargs):
+ super(BiSeNet, self).__init__()
+ self.cp = ContextPath()
+ ## here self.sp is deleted
+ self.ffm = FeatureFusionModule(256, 256)
+ self.conv_out = BiSeNetOutput(256, 256, n_classes)
+ self.conv_out16 = BiSeNetOutput(128, 64, n_classes)
+ self.conv_out32 = BiSeNetOutput(128, 64, n_classes)
+ self.init_weight()
+
+ def forward(self, x):
+ H, W = x.size()[2:]
+ feat_res8, feat_cp8, feat_cp16 = self.cp(x) # here return res3b1 feature
+ feat_sp = feat_res8 # use res3b1 feature to replace spatial path feature
+ feat_fuse = self.ffm(feat_sp, feat_cp8)
+
+ feat_out = self.conv_out(feat_fuse)
+ feat_out16 = self.conv_out16(feat_cp8)
+ feat_out32 = self.conv_out32(feat_cp16)
+
+ feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True)
+ feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True)
+ feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True)
+ return feat_out, feat_out16, feat_out32
+
+ def init_weight(self):
+ for ly in self.children():
+ if isinstance(ly, nn.Conv2d):
+ nn.init.kaiming_normal_(ly.weight, a=1)
+ if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+ def get_params(self):
+ wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []
+ for name, child in self.named_children():
+ child_wd_params, child_nowd_params = child.get_params()
+ if isinstance(child, FeatureFusionModule) or isinstance(child, BiSeNetOutput):
+ lr_mul_wd_params += child_wd_params
+ lr_mul_nowd_params += child_nowd_params
+ else:
+ wd_params += child_wd_params
+ nowd_params += child_nowd_params
+ return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params
+
+
+if __name__ == "__main__":
+ net = BiSeNet(19)
+ net.cuda()
+ net.eval()
+ in_ten = torch.randn(16, 3, 640, 480).cuda()
+ out, out16, out32 = net(in_ten)
+ print(out.shape)
+
+ net.get_params()
diff --git a/SimSwap/parsing_model/resnet.py b/SimSwap/parsing_model/resnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..aa2bf95130e9815ba378cb6f73207068b81a04b9
--- /dev/null
+++ b/SimSwap/parsing_model/resnet.py
@@ -0,0 +1,109 @@
+#!/usr/bin/python
+# -*- encoding: utf-8 -*-
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.model_zoo as modelzoo
+
+# from modules.bn import InPlaceABNSync as BatchNorm2d
+
+resnet18_url = 'https://download.pytorch.org/models/resnet18-5c106cde.pth'
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+ """3x3 convolution with padding"""
+ return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+ padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+ def __init__(self, in_chan, out_chan, stride=1):
+ super(BasicBlock, self).__init__()
+ self.conv1 = conv3x3(in_chan, out_chan, stride)
+ self.bn1 = nn.BatchNorm2d(out_chan)
+ self.conv2 = conv3x3(out_chan, out_chan)
+ self.bn2 = nn.BatchNorm2d(out_chan)
+ self.relu = nn.ReLU(inplace=True)
+ self.downsample = None
+ if in_chan != out_chan or stride != 1:
+ self.downsample = nn.Sequential(
+ nn.Conv2d(in_chan, out_chan,
+ kernel_size=1, stride=stride, bias=False),
+ nn.BatchNorm2d(out_chan),
+ )
+
+ def forward(self, x):
+ residual = self.conv1(x)
+ residual = F.relu(self.bn1(residual))
+ residual = self.conv2(residual)
+ residual = self.bn2(residual)
+
+ shortcut = x
+ if self.downsample is not None:
+ shortcut = self.downsample(x)
+
+ out = shortcut + residual
+ out = self.relu(out)
+ return out
+
+
+def create_layer_basic(in_chan, out_chan, bnum, stride=1):
+ layers = [BasicBlock(in_chan, out_chan, stride=stride)]
+ for i in range(bnum-1):
+ layers.append(BasicBlock(out_chan, out_chan, stride=1))
+ return nn.Sequential(*layers)
+
+
+class Resnet18(nn.Module):
+ def __init__(self):
+ super(Resnet18, self).__init__()
+ self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+ bias=False)
+ self.bn1 = nn.BatchNorm2d(64)
+ self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+ self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1)
+ self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2)
+ self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2)
+ self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2)
+ self.init_weight()
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = F.relu(self.bn1(x))
+ x = self.maxpool(x)
+
+ x = self.layer1(x)
+ feat8 = self.layer2(x) # 1/8
+ feat16 = self.layer3(feat8) # 1/16
+ feat32 = self.layer4(feat16) # 1/32
+ return feat8, feat16, feat32
+
+ def init_weight(self):
+ state_dict = modelzoo.load_url(resnet18_url)
+ self_state_dict = self.state_dict()
+ for k, v in state_dict.items():
+ if 'fc' in k: continue
+ self_state_dict.update({k: v})
+ self.load_state_dict(self_state_dict)
+
+ def get_params(self):
+ wd_params, nowd_params = [], []
+ for name, module in self.named_modules():
+ if isinstance(module, (nn.Linear, nn.Conv2d)):
+ wd_params.append(module.weight)
+ if not module.bias is None:
+ nowd_params.append(module.bias)
+ elif isinstance(module, nn.BatchNorm2d):
+ nowd_params += list(module.parameters())
+ return wd_params, nowd_params
+
+
+if __name__ == "__main__":
+ net = Resnet18()
+ x = torch.randn(16, 3, 224, 224)
+ out = net(x)
+ print(out[0].size())
+ print(out[1].size())
+ print(out[2].size())
+ net.get_params()
diff --git a/SimSwap/pg_modules/blocks.py b/SimSwap/pg_modules/blocks.py
new file mode 100644
index 0000000000000000000000000000000000000000..78bd113bac1cd6486ede92b1ae8d5adfb678eb81
--- /dev/null
+++ b/SimSwap/pg_modules/blocks.py
@@ -0,0 +1,325 @@
+import functools
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import spectral_norm
+
+
+### single layers
+
+
+def conv2d(*args, **kwargs):
+ return spectral_norm(nn.Conv2d(*args, **kwargs))
+
+
+def convTranspose2d(*args, **kwargs):
+ return spectral_norm(nn.ConvTranspose2d(*args, **kwargs))
+
+
+def embedding(*args, **kwargs):
+ return spectral_norm(nn.Embedding(*args, **kwargs))
+
+
+def linear(*args, **kwargs):
+ return spectral_norm(nn.Linear(*args, **kwargs))
+
+
+def NormLayer(c, mode='batch'):
+ if mode == 'group':
+ return nn.GroupNorm(c//2, c)
+ elif mode == 'batch':
+ return nn.BatchNorm2d(c)
+
+
+### Activations
+
+
+class GLU(nn.Module):
+ def forward(self, x):
+ nc = x.size(1)
+ assert nc % 2 == 0, 'channels dont divide 2!'
+ nc = int(nc/2)
+ return x[:, :nc] * torch.sigmoid(x[:, nc:])
+
+
+class Swish(nn.Module):
+ def forward(self, feat):
+ return feat * torch.sigmoid(feat)
+
+
+### Upblocks
+
+
+class InitLayer(nn.Module):
+ def __init__(self, nz, channel, sz=4):
+ super().__init__()
+
+ self.init = nn.Sequential(
+ convTranspose2d(nz, channel*2, sz, 1, 0, bias=False),
+ NormLayer(channel*2),
+ GLU(),
+ )
+
+ def forward(self, noise):
+ noise = noise.view(noise.shape[0], -1, 1, 1)
+ return self.init(noise)
+
+
+def UpBlockSmall(in_planes, out_planes):
+ block = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='nearest'),
+ conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False),
+ NormLayer(out_planes*2), GLU())
+ return block
+
+
+class UpBlockSmallCond(nn.Module):
+ def __init__(self, in_planes, out_planes, z_dim):
+ super().__init__()
+ self.in_planes = in_planes
+ self.out_planes = out_planes
+ self.up = nn.Upsample(scale_factor=2, mode='nearest')
+ self.conv = conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False)
+
+ which_bn = functools.partial(CCBN, which_linear=linear, input_size=z_dim)
+ self.bn = which_bn(2*out_planes)
+ self.act = GLU()
+
+ def forward(self, x, c):
+ x = self.up(x)
+ x = self.conv(x)
+ x = self.bn(x, c)
+ x = self.act(x)
+ return x
+
+
+def UpBlockBig(in_planes, out_planes):
+ block = nn.Sequential(
+ nn.Upsample(scale_factor=2, mode='nearest'),
+ conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False),
+ NoiseInjection(),
+ NormLayer(out_planes*2), GLU(),
+ conv2d(out_planes, out_planes*2, 3, 1, 1, bias=False),
+ NoiseInjection(),
+ NormLayer(out_planes*2), GLU()
+ )
+ return block
+
+
+class UpBlockBigCond(nn.Module):
+ def __init__(self, in_planes, out_planes, z_dim):
+ super().__init__()
+ self.in_planes = in_planes
+ self.out_planes = out_planes
+ self.up = nn.Upsample(scale_factor=2, mode='nearest')
+ self.conv1 = conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False)
+ self.conv2 = conv2d(out_planes, out_planes*2, 3, 1, 1, bias=False)
+
+ which_bn = functools.partial(CCBN, which_linear=linear, input_size=z_dim)
+ self.bn1 = which_bn(2*out_planes)
+ self.bn2 = which_bn(2*out_planes)
+ self.act = GLU()
+ self.noise = NoiseInjection()
+
+ def forward(self, x, c):
+ # block 1
+ x = self.up(x)
+ x = self.conv1(x)
+ x = self.noise(x)
+ x = self.bn1(x, c)
+ x = self.act(x)
+
+ # block 2
+ x = self.conv2(x)
+ x = self.noise(x)
+ x = self.bn2(x, c)
+ x = self.act(x)
+
+ return x
+
+
+class SEBlock(nn.Module):
+ def __init__(self, ch_in, ch_out):
+ super().__init__()
+ self.main = nn.Sequential(
+ nn.AdaptiveAvgPool2d(4),
+ conv2d(ch_in, ch_out, 4, 1, 0, bias=False),
+ Swish(),
+ conv2d(ch_out, ch_out, 1, 1, 0, bias=False),
+ nn.Sigmoid(),
+ )
+
+ def forward(self, feat_small, feat_big):
+ return feat_big * self.main(feat_small)
+
+
+### Downblocks
+
+
+class SeparableConv2d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, bias=False):
+ super(SeparableConv2d, self).__init__()
+ self.depthwise = conv2d(in_channels, in_channels, kernel_size=kernel_size,
+ groups=in_channels, bias=bias, padding=1)
+ self.pointwise = conv2d(in_channels, out_channels,
+ kernel_size=1, bias=bias)
+
+ def forward(self, x):
+ out = self.depthwise(x)
+ out = self.pointwise(out)
+ return out
+
+
+class DownBlock(nn.Module):
+ def __init__(self, in_planes, out_planes, separable=False):
+ super().__init__()
+ if not separable:
+ self.main = nn.Sequential(
+ conv2d(in_planes, out_planes, 4, 2, 1),
+ NormLayer(out_planes),
+ nn.LeakyReLU(0.2, inplace=True),
+ )
+ else:
+ self.main = nn.Sequential(
+ SeparableConv2d(in_planes, out_planes, 3),
+ NormLayer(out_planes),
+ nn.LeakyReLU(0.2, inplace=True),
+ nn.AvgPool2d(2, 2),
+ )
+
+ def forward(self, feat):
+ return self.main(feat)
+
+
+class DownBlockPatch(nn.Module):
+ def __init__(self, in_planes, out_planes, separable=False):
+ super().__init__()
+ self.main = nn.Sequential(
+ DownBlock(in_planes, out_planes, separable),
+ conv2d(out_planes, out_planes, 1, 1, 0, bias=False),
+ NormLayer(out_planes),
+ nn.LeakyReLU(0.2, inplace=True),
+ )
+
+ def forward(self, feat):
+ return self.main(feat)
+
+
+### CSM
+
+
+class ResidualConvUnit(nn.Module):
+ def __init__(self, cin, activation, bn):
+ super().__init__()
+ self.conv = nn.Conv2d(cin, cin, kernel_size=3, stride=1, padding=1, bias=True)
+ self.skip_add = nn.quantized.FloatFunctional()
+
+ def forward(self, x):
+ return self.skip_add.add(self.conv(x), x)
+
+
+class FeatureFusionBlock(nn.Module):
+ def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, lowest=False):
+ super().__init__()
+
+ self.deconv = deconv
+ self.align_corners = align_corners
+
+ self.expand = expand
+ out_features = features
+ if self.expand==True:
+ out_features = features//2
+
+ self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
+ self.skip_add = nn.quantized.FloatFunctional()
+
+ def forward(self, *xs):
+ output = xs[0]
+
+ if len(xs) == 2:
+ output = self.skip_add.add(output, xs[1])
+
+ output = nn.functional.interpolate(
+ output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
+ )
+
+ output = self.out_conv(output)
+
+ return output
+
+
+### Misc
+
+
+class NoiseInjection(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.weight = nn.Parameter(torch.zeros(1), requires_grad=True)
+
+ def forward(self, feat, noise=None):
+ if noise is None:
+ batch, _, height, width = feat.shape
+ noise = torch.randn(batch, 1, height, width).to(feat.device)
+
+ return feat + self.weight * noise
+
+
+class CCBN(nn.Module):
+ ''' conditional batchnorm '''
+ def __init__(self, output_size, input_size, which_linear, eps=1e-5, momentum=0.1):
+ super().__init__()
+ self.output_size, self.input_size = output_size, input_size
+
+ # Prepare gain and bias layers
+ self.gain = which_linear(input_size, output_size)
+ self.bias = which_linear(input_size, output_size)
+
+ # epsilon to avoid dividing by 0
+ self.eps = eps
+ # Momentum
+ self.momentum = momentum
+
+ self.register_buffer('stored_mean', torch.zeros(output_size))
+ self.register_buffer('stored_var', torch.ones(output_size))
+
+ def forward(self, x, y):
+ # Calculate class-conditional gains and biases
+ gain = (1 + self.gain(y)).view(y.size(0), -1, 1, 1)
+ bias = self.bias(y).view(y.size(0), -1, 1, 1)
+ out = F.batch_norm(x, self.stored_mean, self.stored_var, None, None,
+ self.training, 0.1, self.eps)
+ return out * gain + bias
+
+
+class Interpolate(nn.Module):
+ """Interpolation module."""
+
+ def __init__(self, size, mode='bilinear', align_corners=False):
+ """Init.
+ Args:
+ scale_factor (float): scaling
+ mode (str): interpolation mode
+ """
+ super(Interpolate, self).__init__()
+
+ self.interp = nn.functional.interpolate
+ self.size = size
+ self.mode = mode
+ self.align_corners = align_corners
+
+ def forward(self, x):
+ """Forward pass.
+ Args:
+ x (tensor): input
+ Returns:
+ tensor: interpolated data
+ """
+
+ x = self.interp(
+ x,
+ size=self.size,
+ mode=self.mode,
+ align_corners=self.align_corners,
+ )
+
+ return x
diff --git a/SimSwap/pg_modules/diffaug.py b/SimSwap/pg_modules/diffaug.py
new file mode 100644
index 0000000000000000000000000000000000000000..54020be64733f1a92e454b2959f6297f7dcb9e22
--- /dev/null
+++ b/SimSwap/pg_modules/diffaug.py
@@ -0,0 +1,76 @@
+# Differentiable Augmentation for Data-Efficient GAN Training
+# Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu, and Song Han
+# https://arxiv.org/pdf/2006.10738
+
+import torch
+import torch.nn.functional as F
+
+
+def DiffAugment(x, policy='', channels_first=True):
+ if policy:
+ if not channels_first:
+ x = x.permute(0, 3, 1, 2)
+ for p in policy.split(','):
+ for f in AUGMENT_FNS[p]:
+ x = f(x)
+ if not channels_first:
+ x = x.permute(0, 2, 3, 1)
+ x = x.contiguous()
+ return x
+
+
+def rand_brightness(x):
+ x = x + (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) - 0.5)
+ return x
+
+
+def rand_saturation(x):
+ x_mean = x.mean(dim=1, keepdim=True)
+ x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) * 2) + x_mean
+ return x
+
+
+def rand_contrast(x):
+ x_mean = x.mean(dim=[1, 2, 3], keepdim=True)
+ x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) + 0.5) + x_mean
+ return x
+
+
+def rand_translation(x, ratio=0.125):
+ shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+ translation_x = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device)
+ translation_y = torch.randint(-shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device)
+ grid_batch, grid_x, grid_y = torch.meshgrid(
+ torch.arange(x.size(0), dtype=torch.long, device=x.device),
+ torch.arange(x.size(2), dtype=torch.long, device=x.device),
+ torch.arange(x.size(3), dtype=torch.long, device=x.device),
+ )
+ grid_x = torch.clamp(grid_x + translation_x + 1, 0, x.size(2) + 1)
+ grid_y = torch.clamp(grid_y + translation_y + 1, 0, x.size(3) + 1)
+ x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0])
+ x = x_pad.permute(0, 2, 3, 1).contiguous()[grid_batch, grid_x, grid_y].permute(0, 3, 1, 2)
+ return x
+
+
+def rand_cutout(x, ratio=0.2):
+ cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+ offset_x = torch.randint(0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1], device=x.device)
+ offset_y = torch.randint(0, x.size(3) + (1 - cutout_size[1] % 2), size=[x.size(0), 1, 1], device=x.device)
+ grid_batch, grid_x, grid_y = torch.meshgrid(
+ torch.arange(x.size(0), dtype=torch.long, device=x.device),
+ torch.arange(cutout_size[0], dtype=torch.long, device=x.device),
+ torch.arange(cutout_size[1], dtype=torch.long, device=x.device),
+ )
+ grid_x = torch.clamp(grid_x + offset_x - cutout_size[0] // 2, min=0, max=x.size(2) - 1)
+ grid_y = torch.clamp(grid_y + offset_y - cutout_size[1] // 2, min=0, max=x.size(3) - 1)
+ mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device)
+ mask[grid_batch, grid_x, grid_y] = 0
+ x = x * mask.unsqueeze(1)
+ return x
+
+
+AUGMENT_FNS = {
+ 'color': [rand_brightness, rand_saturation, rand_contrast],
+ 'translation': [rand_translation],
+ 'cutout': [rand_cutout],
+}
diff --git a/SimSwap/pg_modules/projected_discriminator.py b/SimSwap/pg_modules/projected_discriminator.py
new file mode 100644
index 0000000000000000000000000000000000000000..d0c879f98eab71c8446c6b1642aaa0718b4a10c5
--- /dev/null
+++ b/SimSwap/pg_modules/projected_discriminator.py
@@ -0,0 +1,191 @@
+from functools import partial
+import numpy as np
+import torch
+import torch.nn as nn
+
+from pg_modules.blocks import DownBlock, DownBlockPatch, conv2d
+from pg_modules.projector import F_RandomProj
+from pg_modules.diffaug import DiffAugment
+
+
+class SingleDisc(nn.Module):
+ def __init__(self, nc=None, ndf=None, start_sz=256, end_sz=8, head=None, separable=False, patch=False):
+ super().__init__()
+ channel_dict = {4: 512, 8: 512, 16: 256, 32: 128, 64: 64, 128: 64,
+ 256: 32, 512: 16, 1024: 8}
+
+ # interpolate for start sz that are not powers of two
+ if start_sz not in channel_dict.keys():
+ sizes = np.array(list(channel_dict.keys()))
+ start_sz = sizes[np.argmin(abs(sizes - start_sz))]
+ self.start_sz = start_sz
+
+ # if given ndf, allocate all layers with the same ndf
+ if ndf is None:
+ nfc = channel_dict
+ else:
+ nfc = {k: ndf for k, v in channel_dict.items()}
+
+ # for feature map discriminators with nfc not in channel_dict
+ # this is the case for the pretrained backbone (midas.pretrained)
+ if nc is not None and head is None:
+ nfc[start_sz] = nc
+
+ layers = []
+
+ # Head if the initial input is the full modality
+ if head:
+ layers += [conv2d(nc, nfc[256], 3, 1, 1, bias=False),
+ nn.LeakyReLU(0.2, inplace=True)]
+
+ # Down Blocks
+ DB = partial(DownBlockPatch, separable=separable) if patch else partial(DownBlock, separable=separable)
+ while start_sz > end_sz:
+ layers.append(DB(nfc[start_sz], nfc[start_sz//2]))
+ start_sz = start_sz // 2
+
+ layers.append(conv2d(nfc[end_sz], 1, 4, 1, 0, bias=False))
+ self.main = nn.Sequential(*layers)
+
+ def forward(self, x, c):
+ return self.main(x)
+
+
+class SingleDiscCond(nn.Module):
+ def __init__(self, nc=None, ndf=None, start_sz=256, end_sz=8, head=None, separable=False, patch=False, c_dim=1000, cmap_dim=64, embedding_dim=128):
+ super().__init__()
+ self.cmap_dim = cmap_dim
+
+ # midas channels
+ channel_dict = {4: 512, 8: 512, 16: 256, 32: 128, 64: 64, 128: 64,
+ 256: 32, 512: 16, 1024: 8}
+
+ # interpolate for start sz that are not powers of two
+ if start_sz not in channel_dict.keys():
+ sizes = np.array(list(channel_dict.keys()))
+ start_sz = sizes[np.argmin(abs(sizes - start_sz))]
+ self.start_sz = start_sz
+
+ # if given ndf, allocate all layers with the same ndf
+ if ndf is None:
+ nfc = channel_dict
+ else:
+ nfc = {k: ndf for k, v in channel_dict.items()}
+
+ # for feature map discriminators with nfc not in channel_dict
+ # this is the case for the pretrained backbone (midas.pretrained)
+ if nc is not None and head is None:
+ nfc[start_sz] = nc
+
+ layers = []
+
+ # Head if the initial input is the full modality
+ if head:
+ layers += [conv2d(nc, nfc[256], 3, 1, 1, bias=False),
+ nn.LeakyReLU(0.2, inplace=True)]
+
+ # Down Blocks
+ DB = partial(DownBlockPatch, separable=separable) if patch else partial(DownBlock, separable=separable)
+ while start_sz > end_sz:
+ layers.append(DB(nfc[start_sz], nfc[start_sz//2]))
+ start_sz = start_sz // 2
+ self.main = nn.Sequential(*layers)
+
+ # additions for conditioning on class information
+ self.cls = conv2d(nfc[end_sz], self.cmap_dim, 4, 1, 0, bias=False)
+ self.embed = nn.Embedding(num_embeddings=c_dim, embedding_dim=embedding_dim)
+ self.embed_proj = nn.Sequential(
+ nn.Linear(self.embed.embedding_dim, self.cmap_dim),
+ nn.LeakyReLU(0.2, inplace=True),
+ )
+
+ def forward(self, x, c):
+ h = self.main(x)
+ out = self.cls(h)
+
+ # conditioning via projection
+ cmap = self.embed_proj(self.embed(c.argmax(1))).unsqueeze(-1).unsqueeze(-1)
+ out = (out * cmap).sum(dim=1, keepdim=True) * (1 / np.sqrt(self.cmap_dim))
+
+ return out
+
+
+class MultiScaleD(nn.Module):
+ def __init__(
+ self,
+ channels,
+ resolutions,
+ num_discs=4,
+ proj_type=2, # 0 = no projection, 1 = cross channel mixing, 2 = cross scale mixing
+ cond=0,
+ separable=False,
+ patch=False,
+ **kwargs,
+ ):
+ super().__init__()
+
+ assert num_discs in [1, 2, 3, 4]
+
+ # the first disc is on the lowest level of the backbone
+ self.disc_in_channels = channels[:num_discs]
+ self.disc_in_res = resolutions[:num_discs]
+ Disc = SingleDiscCond if cond else SingleDisc
+
+ mini_discs = []
+ for i, (cin, res) in enumerate(zip(self.disc_in_channels, self.disc_in_res)):
+ start_sz = res if not patch else 16
+ mini_discs += [str(i), Disc(nc=cin, start_sz=start_sz, end_sz=8, separable=separable, patch=patch)],
+ self.mini_discs = nn.ModuleDict(mini_discs)
+
+ def forward(self, features, c):
+ all_logits = []
+ for k, disc in self.mini_discs.items():
+ res = disc(features[k], c).view(features[k].size(0), -1)
+ all_logits.append(res)
+
+ all_logits = torch.cat(all_logits, dim=1)
+ return all_logits
+
+
+class ProjectedDiscriminator(torch.nn.Module):
+ def __init__(
+ self,
+ diffaug=True,
+ interp224=True,
+ backbone_kwargs={},
+ **kwargs
+ ):
+ super().__init__()
+ self.diffaug = diffaug
+ self.interp224 = interp224
+ self.feature_network = F_RandomProj(**backbone_kwargs)
+ self.discriminator = MultiScaleD(
+ channels=self.feature_network.CHANNELS,
+ resolutions=self.feature_network.RESOLUTIONS,
+ **backbone_kwargs,
+ )
+
+ def train(self, mode=True):
+ self.feature_network = self.feature_network.train(False)
+ self.discriminator = self.discriminator.train(mode)
+ return self
+
+ def eval(self):
+ return self.train(False)
+
+ def get_feature(self, x):
+ features = self.feature_network(x, get_features=True)
+ return features
+
+ def forward(self, x, c):
+ # if self.diffaug:
+ # x = DiffAugment(x, policy='color,translation,cutout')
+
+ # if self.interp224:
+ # x = F.interpolate(x, 224, mode='bilinear', align_corners=False)
+
+ features,backbone_features = self.feature_network(x)
+ logits = self.discriminator(features, c)
+
+ return logits,backbone_features
+
diff --git a/SimSwap/pg_modules/projector.py b/SimSwap/pg_modules/projector.py
new file mode 100644
index 0000000000000000000000000000000000000000..610a48250587521399e41f53b51c579989015b42
--- /dev/null
+++ b/SimSwap/pg_modules/projector.py
@@ -0,0 +1,158 @@
+import torch
+import torch.nn as nn
+import timm
+from pg_modules.blocks import FeatureFusionBlock
+
+
+def _make_scratch_ccm(scratch, in_channels, cout, expand=False):
+ # shapes
+ out_channels = [cout, cout*2, cout*4, cout*8] if expand else [cout]*4
+
+ scratch.layer0_ccm = nn.Conv2d(in_channels[0], out_channels[0], kernel_size=1, stride=1, padding=0, bias=True)
+ scratch.layer1_ccm = nn.Conv2d(in_channels[1], out_channels[1], kernel_size=1, stride=1, padding=0, bias=True)
+ scratch.layer2_ccm = nn.Conv2d(in_channels[2], out_channels[2], kernel_size=1, stride=1, padding=0, bias=True)
+ scratch.layer3_ccm = nn.Conv2d(in_channels[3], out_channels[3], kernel_size=1, stride=1, padding=0, bias=True)
+
+ scratch.CHANNELS = out_channels
+
+ return scratch
+
+
+def _make_scratch_csm(scratch, in_channels, cout, expand):
+ scratch.layer3_csm = FeatureFusionBlock(in_channels[3], nn.ReLU(False), expand=expand, lowest=True)
+ scratch.layer2_csm = FeatureFusionBlock(in_channels[2], nn.ReLU(False), expand=expand)
+ scratch.layer1_csm = FeatureFusionBlock(in_channels[1], nn.ReLU(False), expand=expand)
+ scratch.layer0_csm = FeatureFusionBlock(in_channels[0], nn.ReLU(False))
+
+ # last refinenet does not expand to save channels in higher dimensions
+ scratch.CHANNELS = [cout, cout, cout*2, cout*4] if expand else [cout]*4
+
+ return scratch
+
+
+def _make_efficientnet(model):
+ pretrained = nn.Module()
+ pretrained.layer0 = nn.Sequential(model.conv_stem, model.bn1, model.act1, *model.blocks[0:2])
+ pretrained.layer1 = nn.Sequential(*model.blocks[2:3])
+ pretrained.layer2 = nn.Sequential(*model.blocks[3:5])
+ pretrained.layer3 = nn.Sequential(*model.blocks[5:9])
+ return pretrained
+
+
+def calc_channels(pretrained, inp_res=224):
+ channels = []
+ tmp = torch.zeros(1, 3, inp_res, inp_res)
+
+ # forward pass
+ tmp = pretrained.layer0(tmp)
+ channels.append(tmp.shape[1])
+ tmp = pretrained.layer1(tmp)
+ channels.append(tmp.shape[1])
+ tmp = pretrained.layer2(tmp)
+ channels.append(tmp.shape[1])
+ tmp = pretrained.layer3(tmp)
+ channels.append(tmp.shape[1])
+
+ return channels
+
+
+def _make_projector(im_res, cout, proj_type, expand=False):
+ assert proj_type in [0, 1, 2], "Invalid projection type"
+
+ ### Build pretrained feature network
+ model = timm.create_model('tf_efficientnet_lite0', pretrained=True)
+ pretrained = _make_efficientnet(model)
+
+ # determine resolution of feature maps, this is later used to calculate the number
+ # of down blocks in the discriminators. Interestingly, the best results are achieved
+ # by fixing this to 256, ie., we use the same number of down blocks per discriminator
+ # independent of the dataset resolution
+ im_res = 256
+ pretrained.RESOLUTIONS = [im_res//4, im_res//8, im_res//16, im_res//32]
+ pretrained.CHANNELS = calc_channels(pretrained)
+
+ if proj_type == 0: return pretrained, None
+
+ ### Build CCM
+ scratch = nn.Module()
+ scratch = _make_scratch_ccm(scratch, in_channels=pretrained.CHANNELS, cout=cout, expand=expand)
+ pretrained.CHANNELS = scratch.CHANNELS
+
+ if proj_type == 1: return pretrained, scratch
+
+ ### build CSM
+ scratch = _make_scratch_csm(scratch, in_channels=scratch.CHANNELS, cout=cout, expand=expand)
+
+ # CSM upsamples x2 so the feature map resolution doubles
+ pretrained.RESOLUTIONS = [res*2 for res in pretrained.RESOLUTIONS]
+ pretrained.CHANNELS = scratch.CHANNELS
+
+ return pretrained, scratch
+
+
+class F_RandomProj(nn.Module):
+ def __init__(
+ self,
+ im_res=256,
+ cout=64,
+ expand=True,
+ proj_type=2, # 0 = no projection, 1 = cross channel mixing, 2 = cross scale mixing
+ **kwargs,
+ ):
+ super().__init__()
+ self.proj_type = proj_type
+ self.cout = cout
+ self.expand = expand
+
+ # build pretrained feature network and random decoder (scratch)
+ self.pretrained, self.scratch = _make_projector(im_res=im_res, cout=self.cout, proj_type=self.proj_type, expand=self.expand)
+ self.CHANNELS = self.pretrained.CHANNELS
+ self.RESOLUTIONS = self.pretrained.RESOLUTIONS
+
+ def forward(self, x, get_features=False):
+ # predict feature maps
+ out0 = self.pretrained.layer0(x)
+ out1 = self.pretrained.layer1(out0)
+ out2 = self.pretrained.layer2(out1)
+ out3 = self.pretrained.layer3(out2)
+
+ # start enumerating at the lowest layer (this is where we put the first discriminator)
+ backbone_features = {
+ '0': out0,
+ '1': out1,
+ '2': out2,
+ '3': out3,
+ }
+ if get_features:
+ return backbone_features
+
+ if self.proj_type == 0: return backbone_features
+
+ out0_channel_mixed = self.scratch.layer0_ccm(backbone_features['0'])
+ out1_channel_mixed = self.scratch.layer1_ccm(backbone_features['1'])
+ out2_channel_mixed = self.scratch.layer2_ccm(backbone_features['2'])
+ out3_channel_mixed = self.scratch.layer3_ccm(backbone_features['3'])
+
+ out = {
+ '0': out0_channel_mixed,
+ '1': out1_channel_mixed,
+ '2': out2_channel_mixed,
+ '3': out3_channel_mixed,
+ }
+
+ if self.proj_type == 1: return out
+
+ # from bottom to top
+ out3_scale_mixed = self.scratch.layer3_csm(out3_channel_mixed)
+ out2_scale_mixed = self.scratch.layer2_csm(out3_scale_mixed, out2_channel_mixed)
+ out1_scale_mixed = self.scratch.layer1_csm(out2_scale_mixed, out1_channel_mixed)
+ out0_scale_mixed = self.scratch.layer0_csm(out1_scale_mixed, out0_channel_mixed)
+
+ out = {
+ '0': out0_scale_mixed,
+ '1': out1_scale_mixed,
+ '2': out2_scale_mixed,
+ '3': out3_scale_mixed,
+ }
+
+ return out, backbone_features
diff --git a/SimSwap/predict.py b/SimSwap/predict.py
new file mode 100644
index 0000000000000000000000000000000000000000..3fe60e3b3495b921ad787f11f8e0400090240ab0
--- /dev/null
+++ b/SimSwap/predict.py
@@ -0,0 +1,101 @@
+import cog
+import tempfile
+from pathlib import Path
+import argparse
+import cv2
+import torch
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+from models.models import create_model
+from options.test_options import TestOptions
+from util.reverse2original import reverse2wholeimage
+from util.norm import SpecificNorm
+from test_wholeimage_swapmulti import _totensor
+from insightface_func.face_detect_crop_multi import Face_detect_crop as Face_detect_crop_multi
+from insightface_func.face_detect_crop_single import Face_detect_crop as Face_detect_crop_single
+
+
+class Predictor(cog.Predictor):
+ def setup(self):
+ self.transformer_Arcface = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+ @cog.input("source", type=Path, help="source image")
+ @cog.input("target", type=Path, help="target image")
+ @cog.input("mode", type=str, options=['single', 'all'], default='all',
+ help="swap a single face (the one with highest confidence by face detection) or all faces in the target image")
+ def predict(self, source, target, mode='all'):
+
+ app = Face_detect_crop_multi(name='antelope', root='./insightface_func/models')
+
+ if mode == 'single':
+ app = Face_detect_crop_single(name='antelope', root='./insightface_func/models')
+
+ app.prepare(ctx_id=0, det_thresh=0.6, det_size=(640, 640))
+
+ options = TestOptions()
+ options.initialize()
+ opt = options.parser.parse_args(["--Arc_path", 'arcface_model/arcface_checkpoint.tar', "--pic_a_path", str(source),
+ "--pic_b_path", str(target), "--isTrain", False, "--no_simswaplogo"])
+
+ str_ids = opt.gpu_ids.split(',')
+ opt.gpu_ids = []
+ for str_id in str_ids:
+ id = int(str_id)
+ if id >= 0:
+ opt.gpu_ids.append(id)
+
+ # set gpu ids
+ if len(opt.gpu_ids) > 0:
+ torch.cuda.set_device(opt.gpu_ids[0])
+
+ torch.nn.Module.dump_patches = True
+ model = create_model(opt)
+ model.eval()
+
+ crop_size = opt.crop_size
+ spNorm = SpecificNorm()
+
+ with torch.no_grad():
+ pic_a = opt.pic_a_path
+ img_a_whole = cv2.imread(pic_a)
+ img_a_align_crop, _ = app.get(img_a_whole, crop_size)
+ img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0], cv2.COLOR_BGR2RGB))
+ img_a = self.transformer_Arcface(img_a_align_crop_pil)
+ img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+
+ # convert numpy to tensor
+ img_id = img_id.cuda()
+
+ # create latent id
+ img_id_downsample = F.interpolate(img_id, size=(112,112))
+ latend_id = model.netArc(img_id_downsample)
+ latend_id = F.normalize(latend_id, p=2, dim=1)
+
+ ############## Forward Pass ######################
+
+ pic_b = opt.pic_b_path
+ img_b_whole = cv2.imread(pic_b)
+ img_b_align_crop_list, b_mat_list = app.get(img_b_whole, crop_size)
+
+ swap_result_list = []
+ b_align_crop_tenor_list = []
+
+ for b_align_crop in img_b_align_crop_list:
+ b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop, cv2.COLOR_BGR2RGB))[None, ...].cuda()
+
+ swap_result = model(None, b_align_crop_tenor, latend_id, None, True)[0]
+ swap_result_list.append(swap_result)
+ b_align_crop_tenor_list.append(b_align_crop_tenor)
+
+ net = None
+
+ out_path = Path(tempfile.mkdtemp()) / "output.png"
+
+ reverse2wholeimage(b_align_crop_tenor_list, swap_result_list, b_mat_list, crop_size, img_b_whole, None,
+ str(out_path), opt.no_simswaplogo,
+ pasring_model=net, use_mask=opt.use_mask, norm=spNorm)
+ return out_path
diff --git a/SimSwap/simswaplogo/simswaplogo.png b/SimSwap/simswaplogo/simswaplogo.png
new file mode 100644
index 0000000000000000000000000000000000000000..806d3ac9d6847a8445d50f9d57a5e7f98977aaa4
Binary files /dev/null and b/SimSwap/simswaplogo/simswaplogo.png differ
diff --git a/SimSwap/simswaplogo/socialbook_logo.2020.357eed90add7705e54a8.svg b/SimSwap/simswaplogo/socialbook_logo.2020.357eed90add7705e54a8.svg
new file mode 100644
index 0000000000000000000000000000000000000000..dd04f00b2622f43480553a57e1c45ff8f681e86f
--- /dev/null
+++ b/SimSwap/simswaplogo/socialbook_logo.2020.357eed90add7705e54a8.svg
@@ -0,0 +1,64 @@
+
+
+
+
+
+
+
diff --git a/SimSwap/temp_results/frame_0000000.jpg b/SimSwap/temp_results/frame_0000000.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..a63a81324d7ced06d73bb7ae9290f5d0581d53c7
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000000.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:77c97e6d34545e32bb74e9a464e2bc4d8c01e2e030771a20c03920b7e96db2bc
+size 203105
diff --git a/SimSwap/temp_results/frame_0000001.jpg b/SimSwap/temp_results/frame_0000001.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..a69884c6c9161f392761a8a314bcf3e6dc6c5443
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000001.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:159845fa7b3d000642caa78654de84dd90d14fcf26ab727a4f02e7087ba8b9c8
+size 204103
diff --git a/SimSwap/temp_results/frame_0000002.jpg b/SimSwap/temp_results/frame_0000002.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..c55e1d10b14a63e96db94dadd9c0b47abd54d93c
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000002.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:601344151b1e83eab839aaf4361ac6a15430a4970f42c24923bc22b0e2068690
+size 208084
diff --git a/SimSwap/temp_results/frame_0000003.jpg b/SimSwap/temp_results/frame_0000003.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..95399be96ee52042199d4411e44eeb95b70f0f61
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000003.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:95f8c82946d62799f4f0daea809f6a20aff25971f06e8d3bc44210c96899dce2
+size 212673
diff --git a/SimSwap/temp_results/frame_0000004.jpg b/SimSwap/temp_results/frame_0000004.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..372c1bac56117d3b8b7d0e6b2407cccc3aa344c4
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000004.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af0543018c3162bc736ed395fecbacbb050b660b3683ad7095869c3b2e64b057
+size 215929
diff --git a/SimSwap/temp_results/frame_0000005.jpg b/SimSwap/temp_results/frame_0000005.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..9509f10d8ffcd9f6e8769384970a99b94a82d095
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000005.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c3b4cebbab88ac03bccccb8f0d909c65f8ace5cd73957438ec0fffccb5a41ce
+size 217432
diff --git a/SimSwap/temp_results/frame_0000006.jpg b/SimSwap/temp_results/frame_0000006.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..6bbaf6daae35fd168f2fe24eec3d2f8ed0fb6fa4
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000006.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e10722f43a11a32e2a18a17edeaede49edd243462f42dd69ebb94312599537a
+size 216544
diff --git a/SimSwap/temp_results/frame_0000007.jpg b/SimSwap/temp_results/frame_0000007.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..47ff246b35defa7ebfce6a83aa8eecd0e6fe6f8c
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000007.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e72e4d76bc5420f8d01305f88990e5e5341d77a330eadfd43c899ec145fc7f6
+size 215433
diff --git a/SimSwap/temp_results/frame_0000008.jpg b/SimSwap/temp_results/frame_0000008.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..a14d411762669e1ae91c59790fca284f24ca8264
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000008.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:652814bb27e07ffd0f7d43c091374f44158ac212ca0ee22b03216907d7ef04ee
+size 215208
diff --git a/SimSwap/temp_results/frame_0000009.jpg b/SimSwap/temp_results/frame_0000009.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..a84d3550c6f30b2519ad614dbee37e4242576e84
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000009.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:137ac915252671778d2a17197ff1d9e9e2ff15363fe292106c444b395e249c60
+size 212979
diff --git a/SimSwap/temp_results/frame_0000010.jpg b/SimSwap/temp_results/frame_0000010.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..02fa8e0965748d71dd5bf63794ce6620b099af49
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000010.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0cd560f84b6e4df5164deb3a26d0bc80bba92bee90827e73a35a475ce863940
+size 213894
diff --git a/SimSwap/temp_results/frame_0000011.jpg b/SimSwap/temp_results/frame_0000011.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..11aaed6bfe655d8de1ce850633df35a6f76625b9
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000011.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf7bd1daff673438ec85fe73fb53c91831d3294135a241094ff9c53fbc1de55e
+size 211628
diff --git a/SimSwap/temp_results/frame_0000012.jpg b/SimSwap/temp_results/frame_0000012.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..e912186d18ca1d029886c556d63cb83687fac588
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000012.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f8dcbab8c4863bc96fd79db7956c475af96a7b1408733e0b9ec5bc3ca4025a2
+size 204205
diff --git a/SimSwap/temp_results/frame_0000013.jpg b/SimSwap/temp_results/frame_0000013.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..551f84a074b0cb34c25b130343965e053462c70c
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000013.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65951d09ea142918c34ccb31a83a9a7958b30d73f1c2520d5eb7797573d134a2
+size 210537
diff --git a/SimSwap/temp_results/frame_0000014.jpg b/SimSwap/temp_results/frame_0000014.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..0f561fb71f70dac7201675db0c72bcd3f2f45b65
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000014.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c14e0c6d2c0a0e3bd44b2888b1e27e7ba1df3a0099fb2bcf26baa4642058c9f0
+size 212057
diff --git a/SimSwap/temp_results/frame_0000015.jpg b/SimSwap/temp_results/frame_0000015.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..76e69472bbd596edbfae5a0f3849d71fbc849d2f
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000015.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9537686aae6aef68f29d34de7f615d8863ba871ee8333275c595600b4ffdac55
+size 213155
diff --git a/SimSwap/temp_results/frame_0000016.jpg b/SimSwap/temp_results/frame_0000016.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..31ba85b3a4477732e9fc83b2ead2ca93cd633693
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000016.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f744d6ea4a4782b3a8e62174e7ec684ab38e253f312fb87e3d129e1bff6af114
+size 213460
diff --git a/SimSwap/temp_results/frame_0000017.jpg b/SimSwap/temp_results/frame_0000017.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..08b513a7d327b272b8eeb05a69e1e593dec3965b
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000017.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12ae59017e289aa9913e87ab0f27ae84c77ae745b89739b972f05ae845a90e39
+size 213241
diff --git a/SimSwap/temp_results/frame_0000018.jpg b/SimSwap/temp_results/frame_0000018.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..af20ab88ede4028017603a1f47b6570f7495b8bf
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000018.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42721013b28e1d37d687b5fe14c6737bf747c7a9fa8f90856436b28fe6c76fcb
+size 214736
diff --git a/SimSwap/temp_results/frame_0000019.jpg b/SimSwap/temp_results/frame_0000019.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..921b414af6c17062b9eab1c280e3d65ccd4eff3c
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000019.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e97d5b2f44562ff0517be191021d98bdf12be78f1bc02c90c332ad3c1827948
+size 214176
diff --git a/SimSwap/temp_results/frame_0000020.jpg b/SimSwap/temp_results/frame_0000020.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..9ed7d9751ddb4f00d2744ed1c27a906c4aeb6ec2
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000020.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f21047f5cfdd69a728b00e7217bceedecc45edddb0ccd1b5e72cfa9e60554549
+size 215553
diff --git a/SimSwap/temp_results/frame_0000021.jpg b/SimSwap/temp_results/frame_0000021.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..bda38cac420bd2926220276c67e9694fe651c5b9
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000021.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dcd86ca59b6542b23f6ccd2322f7d099bfdb40e75dff10151fba1894d291b094
+size 214035
diff --git a/SimSwap/temp_results/frame_0000022.jpg b/SimSwap/temp_results/frame_0000022.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..aa5ebeba0b01d796f0b3430896ccbeda37f21249
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000022.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:899bc9911b2c94db8e8533a894fab255add415d8fcdec6600644692cdfd04f65
+size 212609
diff --git a/SimSwap/temp_results/frame_0000023.jpg b/SimSwap/temp_results/frame_0000023.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..1d0956845b9a0b934c385acdcc3ae31004ad1d93
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000023.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb504115fa26a31e253ab03246351deda5964936bab808e63f58b6a5906a85da
+size 211231
diff --git a/SimSwap/temp_results/frame_0000024.jpg b/SimSwap/temp_results/frame_0000024.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..81fbe640dfacb87201d89d394b29dd31ed0e1721
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000024.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e156525cc3718ab84d539a5362382e41dfa8fa1770575bb6063722caebb5354e
+size 200119
diff --git a/SimSwap/temp_results/frame_0000025.jpg b/SimSwap/temp_results/frame_0000025.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..852e975f7c77fcde58f794b66c9f07777077a496
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000025.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5ea81f01e0d181c517d1ddf140635fc75cf3c16af384edd78cca034cd653298c
+size 209119
diff --git a/SimSwap/temp_results/frame_0000026.jpg b/SimSwap/temp_results/frame_0000026.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..b929009071e2b0c870773d74a540b07a9dd93c3f
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000026.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b820b68eed0c7d5a9defedae10ecd5e871deb167b012fe7301d2ba91041bd45
+size 212527
diff --git a/SimSwap/temp_results/frame_0000027.jpg b/SimSwap/temp_results/frame_0000027.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..666a17799fe62fad9806b261a8d2615af2d40619
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000027.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:942d51e1c15501750cff2a137beeb824e79bec45fdf2d2e4cfcb829a30b98424
+size 214170
diff --git a/SimSwap/temp_results/frame_0000028.jpg b/SimSwap/temp_results/frame_0000028.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..712f1922d1f5686c04113f917052d4d03ebe796e
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000028.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4379258f71c55bcedf340cecdb0c888de862eaa4968aed7f4643bbf6e21a173b
+size 214850
diff --git a/SimSwap/temp_results/frame_0000029.jpg b/SimSwap/temp_results/frame_0000029.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..89d274888d8667904c4c835bebebb9bdf1aa29fc
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000029.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd35907145a2a839bfcbe863d0b8bbcb5338fa80b4d438bf81da028a142d94b8
+size 214944
diff --git a/SimSwap/temp_results/frame_0000030.jpg b/SimSwap/temp_results/frame_0000030.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..4eb9a998786eada27010fa5e41b981572c7f3a10
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000030.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4221d77717722d4f2686f1ac8c941d4dedcb706070ffc4d9ccb70940bd1bc5a3
+size 213953
diff --git a/SimSwap/temp_results/frame_0000031.jpg b/SimSwap/temp_results/frame_0000031.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..157b9cc5715c2e8fc9629e3c6f05692d137cdae7
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000031.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:633e4ce6b3cec6fdb218c634bc584c5ce4c6c615900c4130cca60759d08f223c
+size 212542
diff --git a/SimSwap/temp_results/frame_0000032.jpg b/SimSwap/temp_results/frame_0000032.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..2977e02e55de2221cac4aba815597359ec1075a5
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000032.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0623c9d18968d76c92b57812b64fcd54af27025d1400f0f44f9778eba606a8e4
+size 210477
diff --git a/SimSwap/temp_results/frame_0000033.jpg b/SimSwap/temp_results/frame_0000033.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..7332133777319eb2a76da66d9473ab14151ace9a
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000033.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aecfc558afec7c3a0d1bd59ff3d1b23e00214c80cd88a7a7489e703cb0f9091b
+size 211566
diff --git a/SimSwap/temp_results/frame_0000034.jpg b/SimSwap/temp_results/frame_0000034.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..a60f70a4f026e9bf9427940beba3efd17a8d9f52
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000034.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e5f2661046d10e245019f1f7f83e6840e41b28c17b09d0d5a4d96d51320578f
+size 209798
diff --git a/SimSwap/temp_results/frame_0000035.jpg b/SimSwap/temp_results/frame_0000035.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..ca45f4ede486bf0379f4773d5c8f767e230e60ce
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000035.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46df4c2300b417961b055a32cd5616aa8bcfca9cd655c8e024434fc6c339799c
+size 208518
diff --git a/SimSwap/temp_results/frame_0000036.jpg b/SimSwap/temp_results/frame_0000036.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..4bd33dd8ec2be2721e4bd4c7fd4e4f84247e7850
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000036.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e93d387dd7b886f44656c79666e350bea9d4c719ab771d0f49e3f75fcf0f80cf
+size 201139
diff --git a/SimSwap/temp_results/frame_0000037.jpg b/SimSwap/temp_results/frame_0000037.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..33159a95120784b3474cf5d75eec8269873eb700
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000037.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bcc16763ed3760398296b3fbb9a652198eb1f6542725f741faae257c3b26b9e9
+size 208968
diff --git a/SimSwap/temp_results/frame_0000038.jpg b/SimSwap/temp_results/frame_0000038.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..f5adaaa66ba8665d53a064b9b5b62afc650eee8a
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000038.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f5ed3ff99934da4bda9e13602ce5a58800e9b8182de0559b05d70f4cea7c8188
+size 211366
diff --git a/SimSwap/temp_results/frame_0000039.jpg b/SimSwap/temp_results/frame_0000039.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..7beea2a88b268a1326ae3477ead233041b9420eb
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000039.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e9ea9ead7f1617d88295a9ce8921330e8ab9d9a13c88c77479e660388775df79
+size 214140
diff --git a/SimSwap/temp_results/frame_0000040.jpg b/SimSwap/temp_results/frame_0000040.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..3c89677ad076026479f3a40907b2cbe5983a6671
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000040.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4e1491eddb7ea49793305f6a66579cf8dde46d19472065335b7c3f4b3962060
+size 216007
diff --git a/SimSwap/temp_results/frame_0000041.jpg b/SimSwap/temp_results/frame_0000041.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..d6b5dbe8d247c4f8daa73962bb20b5ffad592696
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000041.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e30d8e6cc31fbc705a3c853d97587ce9c54f2a935aad1a92718b6c0311b5f5d8
+size 217840
diff --git a/SimSwap/temp_results/frame_0000042.jpg b/SimSwap/temp_results/frame_0000042.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..7630e59e9d603e6304ffd599440c75c5a4d9cd9d
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000042.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:560e14f87fff5e3717ee525a84d34dd3b780cb8456ec36599f50cf42f31a4962
+size 218211
diff --git a/SimSwap/temp_results/frame_0000043.jpg b/SimSwap/temp_results/frame_0000043.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..37a9250851a8496d788f953344937ad1ccb7594b
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000043.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f8e26b5687336e4cbf933d451cbdf2cdd1635758598cc82b9486970d7a011e8c
+size 218628
diff --git a/SimSwap/temp_results/frame_0000044.jpg b/SimSwap/temp_results/frame_0000044.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..56b59a535abe8497db708f1b46c18ba302b990d0
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000044.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2c5c24cb24dcca322867e286ad064ec1eec6edde2689295d0d39657ca1a7d72
+size 219788
diff --git a/SimSwap/temp_results/frame_0000045.jpg b/SimSwap/temp_results/frame_0000045.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..214793cd52c37ad3df47fca99f4bcee890f39889
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000045.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a47618fbd11e2191b016ab1cd275079ed6f599528679f236ec8ac8b51d91c75
+size 220737
diff --git a/SimSwap/temp_results/frame_0000046.jpg b/SimSwap/temp_results/frame_0000046.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..9cfd6a941f872bacd6d74745f71e62a91ab1a003
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000046.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52d8aee357a56ddb6255cea1533ceb037b695acf6bdef8c756e715275982c60f
+size 221058
diff --git a/SimSwap/temp_results/frame_0000047.jpg b/SimSwap/temp_results/frame_0000047.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..dda3e8e2af40e0b10628b6e3d2ad0f1263b41e5e
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000047.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d726337556b1d4e36f7d966052ed934e7e46cd03107ea2644911f132bc4d816e
+size 217380
diff --git a/SimSwap/temp_results/frame_0000048.jpg b/SimSwap/temp_results/frame_0000048.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..6853585fee8ed816e7e9cfc8331859fe17f2d3b5
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000048.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd4ca042c1bf49b586a06919e8c14b7af63447133fce64b5a4c2fd9bdd360b7c
+size 205828
diff --git a/SimSwap/temp_results/frame_0000049.jpg b/SimSwap/temp_results/frame_0000049.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..be7511541af29371afae1c6688cbe15f1d750e64
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000049.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:22222022a91bd9bbf81839208739a802475c642dd228323fedfd4065a723c5cb
+size 213633
diff --git a/SimSwap/temp_results/frame_0000050.jpg b/SimSwap/temp_results/frame_0000050.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..83499a10b92d043940e9aedf981ce5111869036a
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000050.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:914ddf974ad9a9c8ebc189dbc5ae7c60ec9de20e5cc1fe565a37b76f7137dc44
+size 214044
diff --git a/SimSwap/temp_results/frame_0000051.jpg b/SimSwap/temp_results/frame_0000051.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..caa5b7e7ffe786d403e6ef5e008f06e78d10aa8f
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000051.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a48d3ab99303c7538e37e379a531c6a46fca3a6354b2516934c98308deb7706
+size 214416
diff --git a/SimSwap/temp_results/frame_0000052.jpg b/SimSwap/temp_results/frame_0000052.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..64cf350535d68e9a67d18ae4f7807bc975d1a397
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000052.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3a861a985017be37b166705180128142b1b406b29ab1cd15b115408d4bd4199
+size 214958
diff --git a/SimSwap/temp_results/frame_0000053.jpg b/SimSwap/temp_results/frame_0000053.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..524d4f8afe89b616a646dd802444f2547f920075
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000053.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57726258becf74fba217c21cf1166d4731a561a3ed91789cacd5b35425de1bbf
+size 216580
diff --git a/SimSwap/temp_results/frame_0000054.jpg b/SimSwap/temp_results/frame_0000054.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..658aa4e5eaffb38d07cc6f9a30071e67595c8223
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000054.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f54021c886cc0789687f06ed67e6ff11d33cdaca0c743e5508be6dbc6235794c
+size 215563
diff --git a/SimSwap/temp_results/frame_0000055.jpg b/SimSwap/temp_results/frame_0000055.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..2a3bacfaf52036dff0ada817f68037c8a7dc38e7
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000055.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7701d72766011c4acf54712fdea5a834624817ca1080edc394da45be054c07cc
+size 216696
diff --git a/SimSwap/temp_results/frame_0000056.jpg b/SimSwap/temp_results/frame_0000056.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..db395806022bd7422f20e101c8a36d1335faed65
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000056.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c0bf65a5538a6e24cdea7982a59f3acd5b72eef6a752a70707b57e66c766da77
+size 217600
diff --git a/SimSwap/temp_results/frame_0000057.jpg b/SimSwap/temp_results/frame_0000057.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..cfa851888c3f2bdc25c66d13d400a69b792ba799
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000057.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66b7789ba4f9f2be404138e4b37a7858c0d20c51a467d400d8807e63db5a1ce5
+size 219816
diff --git a/SimSwap/temp_results/frame_0000058.jpg b/SimSwap/temp_results/frame_0000058.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..4e431aeaecd2655c8c449fc49c996ef28a693210
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000058.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f4c34f34071f8426c344c7a9d231c6f1fa199ff4ff3f51580849e8e2ac9a154
+size 218683
diff --git a/SimSwap/temp_results/frame_0000059.jpg b/SimSwap/temp_results/frame_0000059.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..acea385dac1075eba3860f9614f2f70953a6f8dd
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000059.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71d43c70b922e369a746c95be0d511127ba0e9ee272324367377d28d38fa83e7
+size 216792
diff --git a/SimSwap/temp_results/frame_0000060.jpg b/SimSwap/temp_results/frame_0000060.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..c6eccc1986b7185c5c66b9b230f406d2b5547d4e
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000060.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5cd607887c5065eecdb36fdae2363fe4d93ec680be1bfc801aae0ed9ede13292
+size 211876
diff --git a/SimSwap/temp_results/frame_0000061.jpg b/SimSwap/temp_results/frame_0000061.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..5827a2975d072033622e1a3f9bb2e459f47dd19d
--- /dev/null
+++ b/SimSwap/temp_results/frame_0000061.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d4d632efeb86660d92a549c2d46a2a149776ed8e84728f22560bfe19771204d
+size 214794
diff --git a/SimSwap/test_one_image.py b/SimSwap/test_one_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..88e2f7fc33416c952bba2899ab4d4e421b643ee0
--- /dev/null
+++ b/SimSwap/test_one_image.py
@@ -0,0 +1,86 @@
+
+import cv2
+import torch
+import fractions
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+from models.models import create_model
+from options.test_options import TestOptions
+
+
+def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
+
+transformer = transforms.Compose([
+ transforms.ToTensor(),
+ #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+transformer_Arcface = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+detransformer = transforms.Compose([
+ transforms.Normalize([0, 0, 0], [1/0.229, 1/0.224, 1/0.225]),
+ transforms.Normalize([-0.485, -0.456, -0.406], [1, 1, 1])
+ ])
+if __name__ == '__main__':
+ opt = TestOptions().parse()
+
+ start_epoch, epoch_iter = 1, 0
+
+ torch.nn.Module.dump_patches = True
+ model = create_model(opt)
+ model.eval()
+
+ with torch.no_grad():
+
+ pic_a = opt.pic_a_path
+ img_a = Image.open(pic_a).convert('RGB')
+ img_a = transformer_Arcface(img_a)
+ img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+
+ pic_b = opt.pic_b_path
+
+ img_b = Image.open(pic_b).convert('RGB')
+ img_b = transformer(img_b)
+ img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
+
+ # convert numpy to tensor
+ img_id = img_id.cuda()
+ img_att = img_att.cuda()
+
+ #create latent id
+ img_id_downsample = F.interpolate(img_id, size=(112,112))
+ latend_id = model.netArc(img_id_downsample)
+ latend_id = latend_id.detach().to('cpu')
+ latend_id = latend_id/np.linalg.norm(latend_id,axis=1,keepdims=True)
+ latend_id = latend_id.to('cuda')
+
+
+ ############## Forward Pass ######################
+ img_fake = model(img_id, img_att, latend_id, latend_id, True)
+
+
+ for i in range(img_id.shape[0]):
+ if i == 0:
+ row1 = img_id[i]
+ row2 = img_att[i]
+ row3 = img_fake[i]
+ else:
+ row1 = torch.cat([row1, img_id[i]], dim=2)
+ row2 = torch.cat([row2, img_att[i]], dim=2)
+ row3 = torch.cat([row3, img_fake[i]], dim=2)
+
+ #full = torch.cat([row1, row2, row3], dim=1).detach()
+ full = row3.detach()
+ full = full.permute(1, 2, 0)
+ output = full.to('cpu')
+ output = np.array(output)
+ output = output[..., ::-1]
+
+ output = output*255
+
+ cv2.imwrite(opt.output_path + 'result.jpg', output)
\ No newline at end of file
diff --git a/SimSwap/test_video_swap_multispecific.py b/SimSwap/test_video_swap_multispecific.py
new file mode 100644
index 0000000000000000000000000000000000000000..fc9fd71f461040db8509c0ecac770ea0cf8cf94a
--- /dev/null
+++ b/SimSwap/test_video_swap_multispecific.py
@@ -0,0 +1,99 @@
+
+import cv2
+import torch
+import fractions
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+from models.models import create_model
+from options.test_options import TestOptions
+from insightface_func.face_detect_crop_multi import Face_detect_crop
+from util.videoswap_multispecific import video_swap
+import os
+import glob
+
+def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
+
+transformer = transforms.Compose([
+ transforms.ToTensor(),
+ #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+transformer_Arcface = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+# detransformer = transforms.Compose([
+# transforms.Normalize([0, 0, 0], [1/0.229, 1/0.224, 1/0.225]),
+# transforms.Normalize([-0.485, -0.456, -0.406], [1, 1, 1])
+# ])
+
+
+if __name__ == '__main__':
+ opt = TestOptions().parse()
+ pic_specific = opt.pic_specific_path
+ start_epoch, epoch_iter = 1, 0
+ crop_size = opt.crop_size
+
+ multisepcific_dir = opt.multisepcific_dir
+ torch.nn.Module.dump_patches = True
+ if crop_size == 512:
+ opt.which_epoch = 550000
+ opt.name = '512'
+ mode = 'ffhq'
+ else:
+ mode = 'None'
+ model = create_model(opt)
+ model.eval()
+
+
+ app = Face_detect_crop(name='antelope', root='./insightface_func/models')
+ app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode=mode)
+
+ # The specific person to be swapped(source)
+
+ source_specific_id_nonorm_list = []
+ source_path = os.path.join(multisepcific_dir,'SRC_*')
+ source_specific_images_path = sorted(glob.glob(source_path))
+ with torch.no_grad():
+ for source_specific_image_path in source_specific_images_path:
+ specific_person_whole = cv2.imread(source_specific_image_path)
+ specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)
+ specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB))
+ specific_person = transformer_Arcface(specific_person_align_crop_pil)
+ specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape[2])
+ # convert numpy to tensor
+ specific_person = specific_person.cuda()
+ #create latent id
+ specific_person_downsample = F.interpolate(specific_person, size=(112,112))
+ specific_person_id_nonorm = model.netArc(specific_person_downsample)
+ source_specific_id_nonorm_list.append(specific_person_id_nonorm.clone())
+
+
+ # The person who provides id information (list)
+ target_id_norm_list = []
+ target_path = os.path.join(multisepcific_dir,'DST_*')
+ target_images_path = sorted(glob.glob(target_path))
+
+ for target_image_path in target_images_path:
+ img_a_whole = cv2.imread(target_image_path)
+ img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+ img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB))
+ img_a = transformer_Arcface(img_a_align_crop_pil)
+ img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+ # convert numpy to tensor
+ img_id = img_id.cuda()
+ #create latent id
+ img_id_downsample = F.interpolate(img_id, size=(112,112))
+ latend_id = model.netArc(img_id_downsample)
+ latend_id = F.normalize(latend_id, p=2, dim=1)
+ target_id_norm_list.append(latend_id.clone())
+
+ assert len(target_id_norm_list) == len(source_specific_id_nonorm_list), "The number of images in source and target directory must be same !!!"
+
+
+
+ video_swap(opt.video_path, target_id_norm_list,source_specific_id_nonorm_list, opt.id_thres, \
+ model, app, opt.output_path,temp_results_dir=opt.temp_path,no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask,crop_size=crop_size)
+
diff --git a/SimSwap/test_video_swapmulti.py b/SimSwap/test_video_swapmulti.py
new file mode 100644
index 0000000000000000000000000000000000000000..37d77511000c492fdf8808b15de165d85b33c120
--- /dev/null
+++ b/SimSwap/test_video_swapmulti.py
@@ -0,0 +1,88 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 17:03:58
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-24 19:00:34
+Description:
+'''
+
+import cv2
+import torch
+import fractions
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+from models.models import create_model
+from options.test_options import TestOptions
+from insightface_func.face_detect_crop_multi import Face_detect_crop
+from util.videoswap import video_swap
+import os
+
+def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
+
+transformer = transforms.Compose([
+ transforms.ToTensor(),
+ #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+transformer_Arcface = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+# detransformer = transforms.Compose([
+# transforms.Normalize([0, 0, 0], [1/0.229, 1/0.224, 1/0.225]),
+# transforms.Normalize([-0.485, -0.456, -0.406], [1, 1, 1])
+# ])
+
+
+if __name__ == '__main__':
+ opt = TestOptions().parse()
+
+ start_epoch, epoch_iter = 1, 0
+ crop_size = opt.crop_size
+
+ torch.nn.Module.dump_patches = True
+
+ if crop_size == 512:
+ opt.which_epoch = 550000
+ opt.name = '512'
+ mode = 'ffhq'
+ else:
+ mode = 'None'
+ model = create_model(opt)
+ model.eval()
+
+ app = Face_detect_crop(name='antelope', root='./insightface_func/models')
+ app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode = mode)
+
+ with torch.no_grad():
+ pic_a = opt.pic_a_path
+ # img_a = Image.open(pic_a).convert('RGB')
+ img_a_whole = cv2.imread(pic_a)
+ img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+ img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB))
+ img_a = transformer_Arcface(img_a_align_crop_pil)
+ img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+
+ # pic_b = opt.pic_b_path
+ # img_b_whole = cv2.imread(pic_b)
+ # img_b_align_crop, b_mat = app.get(img_b_whole,crop_size)
+ # img_b_align_crop_pil = Image.fromarray(cv2.cvtColor(img_b_align_crop,cv2.COLOR_BGR2RGB))
+ # img_b = transformer(img_b_align_crop_pil)
+ # img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
+
+ # convert numpy to tensor
+ img_id = img_id.cuda()
+ # img_att = img_att.cuda()
+
+ #create latent id
+ img_id_downsample = F.interpolate(img_id, size=(112,112))
+ latend_id = model.netArc(img_id_downsample)
+ latend_id = F.normalize(latend_id, p=2, dim=1)
+
+ video_swap(opt.video_path, latend_id, model, app, opt.output_path,temp_results_dir=opt.temp_path,\
+ no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask,crop_size=crop_size)
+
diff --git a/SimSwap/test_video_swapsingle.py b/SimSwap/test_video_swapsingle.py
new file mode 100644
index 0000000000000000000000000000000000000000..789eafe53c553be40305b9ee1164262ee63e7a7f
--- /dev/null
+++ b/SimSwap/test_video_swapsingle.py
@@ -0,0 +1,89 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 17:03:58
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-24 19:00:38
+Description:
+'''
+
+import cv2
+import torch
+import fractions
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+from models.models import create_model
+from options.test_options import TestOptions
+from insightface_func.face_detect_crop_single import Face_detect_crop
+from util.videoswap import video_swap
+import os
+
+def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
+
+transformer = transforms.Compose([
+ transforms.ToTensor(),
+ #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+transformer_Arcface = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+# detransformer = transforms.Compose([
+# transforms.Normalize([0, 0, 0], [1/0.229, 1/0.224, 1/0.225]),
+# transforms.Normalize([-0.485, -0.456, -0.406], [1, 1, 1])
+# ])
+
+
+if __name__ == '__main__':
+ opt = TestOptions().parse()
+
+ start_epoch, epoch_iter = 1, 0
+ crop_size = opt.crop_size
+
+ torch.nn.Module.dump_patches = True
+ if crop_size == 512:
+ opt.which_epoch = 550000
+ opt.name = '512'
+ mode = 'ffhq'
+ else:
+ mode = 'None'
+ model = create_model(opt)
+ model.eval()
+
+
+ app = Face_detect_crop(name='antelope', root='./insightface_func/models')
+ app.prepare(ctx_id= 0, det_thresh=0.2, det_size=(640,640),mode=mode)
+ with torch.no_grad():
+ pic_a = opt.pic_a_path
+ # img_a = Image.open(pic_a).convert('RGB')
+ img_a_whole = cv2.imread(pic_a)
+ img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+ img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB))
+ img_a = transformer_Arcface(img_a_align_crop_pil)
+ img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+
+
+
+ # pic_b = opt.pic_b_path
+ # img_b_whole = cv2.imread(pic_b)
+ # img_b_align_crop, b_mat = app.get(img_b_whole,crop_size)
+ # img_b_align_crop_pil = Image.fromarray(cv2.cvtColor(img_b_align_crop,cv2.COLOR_BGR2RGB))
+ # img_b = transformer(img_b_align_crop_pil)
+ # img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
+
+ # convert numpy to tensor
+ img_id = img_id.cuda()
+ # img_att = img_att.cuda()
+
+ #create latent id
+ img_id_downsample = F.interpolate(img_id, size=(112,112))
+ latend_id = model.netArc(img_id_downsample)
+ latend_id = F.normalize(latend_id, p=2, dim=1)
+
+ video_swap(opt.video_path, latend_id, model, app, opt.output_path,temp_results_dir=opt.temp_path,\
+ no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask,crop_size=crop_size)
+
diff --git a/SimSwap/test_video_swapspecific.py b/SimSwap/test_video_swapspecific.py
new file mode 100644
index 0000000000000000000000000000000000000000..3af643d2e76e9b258da44e0a535f7148a06aa433
--- /dev/null
+++ b/SimSwap/test_video_swapspecific.py
@@ -0,0 +1,98 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 17:03:58
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-24 19:00:42
+Description:
+'''
+
+import cv2
+import torch
+import fractions
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+from models.models import create_model
+from options.test_options import TestOptions
+from insightface_func.face_detect_crop_multi import Face_detect_crop
+from util.videoswap_specific import video_swap
+import os
+
+def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
+
+transformer = transforms.Compose([
+ transforms.ToTensor(),
+ #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+transformer_Arcface = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+# detransformer = transforms.Compose([
+# transforms.Normalize([0, 0, 0], [1/0.229, 1/0.224, 1/0.225]),
+# transforms.Normalize([-0.485, -0.456, -0.406], [1, 1, 1])
+# ])
+
+
+if __name__ == '__main__':
+ opt = TestOptions().parse()
+ pic_specific = opt.pic_specific_path
+ start_epoch, epoch_iter = 1, 0
+ crop_size = opt.crop_size
+
+ torch.nn.Module.dump_patches = True
+ if crop_size == 512:
+ opt.which_epoch = 550000
+ opt.name = '512'
+ mode = 'ffhq'
+ else:
+ mode = 'None'
+ model = create_model(opt)
+ model.eval()
+
+
+ app = Face_detect_crop(name='antelope', root='./insightface_func/models')
+ app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode=mode)
+ with torch.no_grad():
+ pic_a = opt.pic_a_path
+ # img_a = Image.open(pic_a).convert('RGB')
+ img_a_whole = cv2.imread(pic_a)
+ img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+ img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB))
+ img_a = transformer_Arcface(img_a_align_crop_pil)
+ img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+
+ # pic_b = opt.pic_b_path
+ # img_b_whole = cv2.imread(pic_b)
+ # img_b_align_crop, b_mat = app.get(img_b_whole,crop_size)
+ # img_b_align_crop_pil = Image.fromarray(cv2.cvtColor(img_b_align_crop,cv2.COLOR_BGR2RGB))
+ # img_b = transformer(img_b_align_crop_pil)
+ # img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
+
+ # convert numpy to tensor
+ img_id = img_id.cuda()
+ # img_att = img_att.cuda()
+
+ #create latent id
+ img_id_downsample = F.interpolate(img_id, size=(112,112))
+ latend_id = model.netArc(img_id_downsample)
+ latend_id = F.normalize(latend_id, p=2, dim=1)
+
+
+ # The specific person to be swapped
+ specific_person_whole = cv2.imread(pic_specific)
+ specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)
+ specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB))
+ specific_person = transformer_Arcface(specific_person_align_crop_pil)
+ specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape[2])
+ specific_person = specific_person.cuda()
+ specific_person_downsample = F.interpolate(specific_person, size=(112,112))
+ specific_person_id_nonorm = model.netArc(specific_person_downsample)
+
+ video_swap(opt.video_path, latend_id,specific_person_id_nonorm, opt.id_thres, \
+ model, app, opt.output_path,temp_results_dir=opt.temp_path,no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask,crop_size=crop_size)
+
diff --git a/SimSwap/test_wholeimage_swap_multispecific.py b/SimSwap/test_wholeimage_swap_multispecific.py
new file mode 100644
index 0000000000000000000000000000000000000000..dca31affb303cef1e0e1bbb9b9bec8d926b3f84b
--- /dev/null
+++ b/SimSwap/test_wholeimage_swap_multispecific.py
@@ -0,0 +1,176 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 17:03:58
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-24 19:19:22
+Description:
+'''
+
+import cv2
+import torch
+import fractions
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+from models.models import create_model
+from options.test_options import TestOptions
+from insightface_func.face_detect_crop_multi import Face_detect_crop
+from util.reverse2original import reverse2wholeimage
+import os
+from util.add_watermark import watermark_image
+import torch.nn as nn
+from util.norm import SpecificNorm
+import glob
+from parsing_model.model import BiSeNet
+
+def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
+
+transformer_Arcface = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+def _totensor(array):
+ tensor = torch.from_numpy(array)
+ img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
+ return img.float().div(255)
+
+def _toarctensor(array):
+ tensor = torch.from_numpy(array)
+ img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
+ return img.float().div(255)
+
+if __name__ == '__main__':
+ opt = TestOptions().parse()
+
+ start_epoch, epoch_iter = 1, 0
+ crop_size = opt.crop_size
+
+ multisepcific_dir = opt.multisepcific_dir
+
+ torch.nn.Module.dump_patches = True
+
+ if crop_size == 512:
+ opt.which_epoch = 550000
+ opt.name = '512'
+ mode = 'ffhq'
+ else:
+ mode = 'None'
+
+ logoclass = watermark_image('./simswaplogo/simswaplogo.png')
+ model = create_model(opt)
+ model.eval()
+ mse = torch.nn.MSELoss().cuda()
+
+ spNorm =SpecificNorm()
+
+
+ app = Face_detect_crop(name='antelope', root='./insightface_func/models')
+ app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode = mode)
+
+ with torch.no_grad():
+ # The specific person to be swapped(source)
+
+ source_specific_id_nonorm_list = []
+ source_path = os.path.join(multisepcific_dir,'SRC_*')
+ source_specific_images_path = sorted(glob.glob(source_path))
+
+ for source_specific_image_path in source_specific_images_path:
+ specific_person_whole = cv2.imread(source_specific_image_path)
+ specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)
+ specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB))
+ specific_person = transformer_Arcface(specific_person_align_crop_pil)
+ specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape[2])
+ # convert numpy to tensor
+ specific_person = specific_person.cuda()
+ #create latent id
+ specific_person_downsample = F.interpolate(specific_person, size=(112,112))
+ specific_person_id_nonorm = model.netArc(specific_person_downsample)
+ source_specific_id_nonorm_list.append(specific_person_id_nonorm.clone())
+
+
+ # The person who provides id information (list)
+ target_id_norm_list = []
+ target_path = os.path.join(multisepcific_dir,'DST_*')
+ target_images_path = sorted(glob.glob(target_path))
+
+ for target_image_path in target_images_path:
+ img_a_whole = cv2.imread(target_image_path)
+ img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+ img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB))
+ img_a = transformer_Arcface(img_a_align_crop_pil)
+ img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+ # convert numpy to tensor
+ img_id = img_id.cuda()
+ #create latent id
+ img_id_downsample = F.interpolate(img_id, size=(112,112))
+ latend_id = model.netArc(img_id_downsample)
+ latend_id = F.normalize(latend_id, p=2, dim=1)
+ target_id_norm_list.append(latend_id.clone())
+
+ assert len(target_id_norm_list) == len(source_specific_id_nonorm_list), "The number of images in source and target directory must be same !!!"
+
+ ############## Forward Pass ######################
+
+ pic_b = opt.pic_b_path
+ img_b_whole = cv2.imread(pic_b)
+
+ img_b_align_crop_list, b_mat_list = app.get(img_b_whole,crop_size)
+ # detect_results = None
+ swap_result_list = []
+
+ id_compare_values = []
+ b_align_crop_tenor_list = []
+ for b_align_crop in img_b_align_crop_list:
+
+ b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
+
+ b_align_crop_tenor_arcnorm = spNorm(b_align_crop_tenor)
+ b_align_crop_tenor_arcnorm_downsample = F.interpolate(b_align_crop_tenor_arcnorm, size=(112,112))
+ b_align_crop_id_nonorm = model.netArc(b_align_crop_tenor_arcnorm_downsample)
+
+ id_compare_values.append([])
+ for source_specific_id_nonorm_tmp in source_specific_id_nonorm_list:
+ id_compare_values[-1].append(mse(b_align_crop_id_nonorm,source_specific_id_nonorm_tmp).detach().cpu().numpy())
+ b_align_crop_tenor_list.append(b_align_crop_tenor)
+
+ id_compare_values_array = np.array(id_compare_values).transpose(1,0)
+ min_indexs = np.argmin(id_compare_values_array,axis=0)
+ min_value = np.min(id_compare_values_array,axis=0)
+
+ swap_result_list = []
+ swap_result_matrix_list = []
+ swap_result_ori_pic_list = []
+
+ for tmp_index, min_index in enumerate(min_indexs):
+ if min_value[tmp_index] < opt.id_thres:
+ swap_result = model(None, b_align_crop_tenor_list[tmp_index], target_id_norm_list[min_index], None, True)[0]
+ swap_result_list.append(swap_result)
+ swap_result_matrix_list.append(b_mat_list[tmp_index])
+ swap_result_ori_pic_list.append(b_align_crop_tenor_list[tmp_index])
+ else:
+ pass
+
+ if len(swap_result_list) !=0:
+
+ if opt.use_mask:
+ n_classes = 19
+ net = BiSeNet(n_classes=n_classes)
+ net.cuda()
+ save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+ net.load_state_dict(torch.load(save_pth))
+ net.eval()
+ else:
+ net =None
+
+ reverse2wholeimage(swap_result_ori_pic_list, swap_result_list, swap_result_matrix_list, crop_size, img_b_whole, logoclass,\
+ os.path.join(opt.output_path, 'result_whole_swap_multispecific.jpg'), opt.no_simswaplogo,pasring_model =net,use_mask=opt.use_mask, norm = spNorm)
+
+ print(' ')
+
+ print('************ Done ! ************')
+
+ else:
+ print('The people you specified are not found on the picture: {}'.format(pic_b))
diff --git a/SimSwap/test_wholeimage_swapmulti.py b/SimSwap/test_wholeimage_swapmulti.py
new file mode 100644
index 0000000000000000000000000000000000000000..19380578c66eb8ff4a7db40980a948269b3c30bb
--- /dev/null
+++ b/SimSwap/test_wholeimage_swapmulti.py
@@ -0,0 +1,110 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 17:03:58
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-24 19:19:26
+Description:
+'''
+
+import cv2
+import torch
+import fractions
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+from models.models import create_model
+from options.test_options import TestOptions
+from insightface_func.face_detect_crop_multi import Face_detect_crop
+from util.reverse2original import reverse2wholeimage
+import os
+from util.add_watermark import watermark_image
+from util.norm import SpecificNorm
+from parsing_model.model import BiSeNet
+
+def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
+
+transformer_Arcface = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+def _totensor(array):
+ tensor = torch.from_numpy(array)
+ img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
+ return img.float().div(255)
+
+if __name__ == '__main__':
+ opt = TestOptions().parse()
+
+ start_epoch, epoch_iter = 1, 0
+ crop_size = opt.crop_size
+
+ torch.nn.Module.dump_patches = True
+ if crop_size == 512:
+ opt.which_epoch = 550000
+ opt.name = '512'
+ mode = 'ffhq'
+ else:
+ mode = 'None'
+ logoclass = watermark_image('./simswaplogo/simswaplogo.png')
+ model = create_model(opt)
+ model.eval()
+ spNorm =SpecificNorm()
+
+ app = Face_detect_crop(name='antelope', root='./insightface_func/models')
+ app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode=mode)
+
+ with torch.no_grad():
+ pic_a = opt.pic_a_path
+
+ img_a_whole = cv2.imread(pic_a)
+ img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+ img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB))
+ img_a = transformer_Arcface(img_a_align_crop_pil)
+ img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+
+ # convert numpy to tensor
+ img_id = img_id.cuda()
+
+ #create latent id
+ img_id_downsample = F.interpolate(img_id, size=(112,112))
+ latend_id = model.netArc(img_id_downsample)
+ latend_id = F.normalize(latend_id, p=2, dim=1)
+
+
+ ############## Forward Pass ######################
+
+ pic_b = opt.pic_b_path
+ img_b_whole = cv2.imread(pic_b)
+
+ img_b_align_crop_list, b_mat_list = app.get(img_b_whole,crop_size)
+ # detect_results = None
+ swap_result_list = []
+ b_align_crop_tenor_list = []
+
+ for b_align_crop in img_b_align_crop_list:
+
+ b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
+
+ swap_result = model(None, b_align_crop_tenor, latend_id, None, True)[0]
+ swap_result_list.append(swap_result)
+ b_align_crop_tenor_list.append(b_align_crop_tenor)
+
+
+ if opt.use_mask:
+ n_classes = 19
+ net = BiSeNet(n_classes=n_classes)
+ net.cuda()
+ save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+ net.load_state_dict(torch.load(save_pth))
+ net.eval()
+ else:
+ net =None
+
+ reverse2wholeimage(b_align_crop_tenor_list,swap_result_list, b_mat_list, crop_size, img_b_whole, logoclass, \
+ os.path.join(opt.output_path, 'result_whole_swapmulti.jpg'),opt.no_simswaplogo,pasring_model =net,use_mask=opt.use_mask, norm = spNorm)
+ print(' ')
+
+ print('************ Done ! ************')
diff --git a/SimSwap/test_wholeimage_swapsingle.py b/SimSwap/test_wholeimage_swapsingle.py
new file mode 100644
index 0000000000000000000000000000000000000000..e35d8437f8a256b5bdf3eca9b201131928a6d5cd
--- /dev/null
+++ b/SimSwap/test_wholeimage_swapsingle.py
@@ -0,0 +1,110 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 17:03:58
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-24 19:19:43
+Description:
+'''
+
+import cv2
+import torch
+import fractions
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+from models.models import create_model
+from options.test_options import TestOptions
+from insightface_func.face_detect_crop_single import Face_detect_crop
+from util.reverse2original import reverse2wholeimage
+import os
+from util.add_watermark import watermark_image
+from util.norm import SpecificNorm
+from parsing_model.model import BiSeNet
+
+def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
+
+transformer_Arcface = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+def _totensor(array):
+ tensor = torch.from_numpy(array)
+ img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
+ return img.float().div(255)
+if __name__ == '__main__':
+ opt = TestOptions().parse()
+
+ start_epoch, epoch_iter = 1, 0
+ crop_size = opt.crop_size
+
+ torch.nn.Module.dump_patches = True
+ if crop_size == 512:
+ opt.which_epoch = 550000
+ opt.name = '512'
+ mode = 'ffhq'
+ else:
+ mode = 'None'
+ logoclass = watermark_image('./simswaplogo/simswaplogo.png')
+ model = create_model(opt)
+ model.eval()
+
+ spNorm =SpecificNorm()
+ app = Face_detect_crop(name='antelope', root='./insightface_func/models')
+ app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode=mode)
+
+ with torch.no_grad():
+ pic_a = opt.pic_a_path
+
+ img_a_whole = cv2.imread(pic_a)
+ img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+ img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB))
+ img_a = transformer_Arcface(img_a_align_crop_pil)
+ img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+
+ # convert numpy to tensor
+ img_id = img_id.cuda()
+
+ #create latent id
+ img_id_downsample = F.interpolate(img_id, size=(112,112))
+ latend_id = model.netArc(img_id_downsample)
+ latend_id = F.normalize(latend_id, p=2, dim=1)
+
+
+ ############## Forward Pass ######################
+
+ pic_b = opt.pic_b_path
+ img_b_whole = cv2.imread(pic_b)
+
+ img_b_align_crop_list, b_mat_list = app.get(img_b_whole,crop_size)
+ # detect_results = None
+ swap_result_list = []
+
+ b_align_crop_tenor_list = []
+
+ for b_align_crop in img_b_align_crop_list:
+
+ b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
+
+ swap_result = model(None, b_align_crop_tenor, latend_id, None, True)[0]
+ swap_result_list.append(swap_result)
+ b_align_crop_tenor_list.append(b_align_crop_tenor)
+
+ if opt.use_mask:
+ n_classes = 19
+ net = BiSeNet(n_classes=n_classes)
+ net.cuda()
+ save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+ net.load_state_dict(torch.load(save_pth))
+ net.eval()
+ else:
+ net =None
+
+ reverse2wholeimage(b_align_crop_tenor_list, swap_result_list, b_mat_list, crop_size, img_b_whole, logoclass, \
+ os.path.join(opt.output_path, 'result_whole_swapsingle.jpg'), opt.no_simswaplogo,pasring_model =net,use_mask=opt.use_mask, norm = spNorm)
+
+ print(' ')
+
+ print('************ Done ! ************')
diff --git a/SimSwap/test_wholeimage_swapspecific.py b/SimSwap/test_wholeimage_swapspecific.py
new file mode 100644
index 0000000000000000000000000000000000000000..b9167f8d387e20d00441a6dc2bcdb1c25b48242d
--- /dev/null
+++ b/SimSwap/test_wholeimage_swapspecific.py
@@ -0,0 +1,150 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 17:03:58
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-24 19:19:47
+Description:
+'''
+
+import cv2
+import torch
+import fractions
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+from models.models import create_model
+from options.test_options import TestOptions
+from insightface_func.face_detect_crop_multi import Face_detect_crop
+from util.reverse2original import reverse2wholeimage
+import os
+from util.add_watermark import watermark_image
+import torch.nn as nn
+from util.norm import SpecificNorm
+from parsing_model.model import BiSeNet
+
+def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
+
+transformer_Arcface = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+def _totensor(array):
+ tensor = torch.from_numpy(array)
+ img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
+ return img.float().div(255)
+
+def _toarctensor(array):
+ tensor = torch.from_numpy(array)
+ img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
+ return img.float().div(255)
+
+if __name__ == '__main__':
+ opt = TestOptions().parse()
+
+ start_epoch, epoch_iter = 1, 0
+ crop_size = opt.crop_size
+
+ torch.nn.Module.dump_patches = True
+ if crop_size == 512:
+ opt.which_epoch = 550000
+ opt.name = '512'
+ mode = 'ffhq'
+ else:
+ mode = 'None'
+ logoclass = watermark_image('./simswaplogo/simswaplogo.png')
+ model = create_model(opt)
+ model.eval()
+ mse = torch.nn.MSELoss().cuda()
+
+ spNorm =SpecificNorm()
+
+
+ app = Face_detect_crop(name='antelope', root='./insightface_func/models')
+ app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640),mode=mode)
+
+ pic_a = opt.pic_a_path
+ pic_specific = opt.pic_specific_path
+
+ # The person who provides id information
+ img_a_whole = cv2.imread(pic_a)
+ img_a_align_crop, _ = app.get(img_a_whole,crop_size)
+ img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB))
+ img_a = transformer_Arcface(img_a_align_crop_pil)
+ img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
+
+ # convert numpy to tensor
+ img_id = img_id.cuda()
+
+ #create latent id
+ img_id_downsample = F.interpolate(img_id, size=(112,112))
+ latend_id = model.netArc(img_id_downsample)
+ latend_id = F.normalize(latend_id, p=2, dim=1)
+
+
+ # The specific person to be swapped
+ specific_person_whole = cv2.imread(pic_specific)
+ specific_person_align_crop, _ = app.get(specific_person_whole,crop_size)
+ specific_person_align_crop_pil = Image.fromarray(cv2.cvtColor(specific_person_align_crop[0],cv2.COLOR_BGR2RGB))
+ specific_person = transformer_Arcface(specific_person_align_crop_pil)
+ specific_person = specific_person.view(-1, specific_person.shape[0], specific_person.shape[1], specific_person.shape[2])
+
+ # convert numpy to tensor
+ specific_person = specific_person.cuda()
+
+ #create latent id
+ specific_person_downsample = F.interpolate(specific_person, size=(112,112))
+ specific_person_id_nonorm = model.netArc(specific_person_downsample)
+ # specific_person_id_norm = F.normalize(specific_person_id_nonorm, p=2, dim=1)
+
+ ############## Forward Pass ######################
+
+ pic_b = opt.pic_b_path
+ img_b_whole = cv2.imread(pic_b)
+
+ img_b_align_crop_list, b_mat_list = app.get(img_b_whole,crop_size)
+ # detect_results = None
+ swap_result_list = []
+
+ id_compare_values = []
+ b_align_crop_tenor_list = []
+ for b_align_crop in img_b_align_crop_list:
+
+ b_align_crop_tenor = _totensor(cv2.cvtColor(b_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
+
+ b_align_crop_tenor_arcnorm = spNorm(b_align_crop_tenor)
+ b_align_crop_tenor_arcnorm_downsample = F.interpolate(b_align_crop_tenor_arcnorm, size=(112,112))
+ b_align_crop_id_nonorm = model.netArc(b_align_crop_tenor_arcnorm_downsample)
+
+ id_compare_values.append(mse(b_align_crop_id_nonorm,specific_person_id_nonorm).detach().cpu().numpy())
+ b_align_crop_tenor_list.append(b_align_crop_tenor)
+
+ id_compare_values_array = np.array(id_compare_values)
+ min_index = np.argmin(id_compare_values_array)
+ min_value = id_compare_values_array[min_index]
+
+ if opt.use_mask:
+ n_classes = 19
+ net = BiSeNet(n_classes=n_classes)
+ net.cuda()
+ save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+ net.load_state_dict(torch.load(save_pth))
+ net.eval()
+ else:
+ net =None
+
+ if min_value < opt.id_thres:
+
+ swap_result = model(None, b_align_crop_tenor_list[min_index], latend_id, None, True)[0]
+
+ reverse2wholeimage([b_align_crop_tenor_list[min_index]], [swap_result], [b_mat_list[min_index]], crop_size, img_b_whole, logoclass, \
+ os.path.join(opt.output_path, 'result_whole_swapspecific.jpg'), opt.no_simswaplogo,pasring_model =net,use_mask=opt.use_mask, norm = spNorm)
+
+ print(' ')
+
+ print('************ Done ! ************')
+
+ else:
+ print('The person you specified is not found on the picture: {}'.format(pic_b))
diff --git a/SimSwap/train.ipynb b/SimSwap/train.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..9dda892309a7b934d3f4f62164c859d0512c8cbb
--- /dev/null
+++ b/SimSwap/train.ipynb
@@ -0,0 +1,192 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "fC7QoKePuJWu"
+ },
+ "source": [
+ "#Training Demo\n",
+ "This is a simple example for training the SimSwap 224*224 with VGGFace2-224.\n",
+ "\n",
+ "Code path: https://github.com/neuralchen/SimSwap\n",
+ "If you like the SimSwap project, please star it!\n",
+ "Paper path: https://arxiv.org/pdf/2106.06340v1.pdf or https://dl.acm.org/doi/10.1145/3394171.3413630"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {
+ "id": "J8WrNaQHuUGC"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Fri Apr 22 12:19:42 2022 \n",
+ "+-----------------------------------------------------------------------------+\n",
+ "| NVIDIA-SMI 456.71 Driver Version: 456.71 CUDA Version: 11.1 |\n",
+ "|-------------------------------+----------------------+----------------------+\n",
+ "| GPU Name TCC/WDDM | Bus-Id Disp.A | Volatile Uncorr. ECC |\n",
+ "| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |\n",
+ "|===============================+======================+======================|\n",
+ "| 0 TITAN Xp WDDM | 00000000:01:00.0 On | N/A |\n",
+ "| 23% 36C P8 15W / 250W | 1135MiB / 12288MiB | 4% Default |\n",
+ "+-------------------------------+----------------------+----------------------+\n",
+ " \n",
+ "+-----------------------------------------------------------------------------+\n",
+ "| Processes: |\n",
+ "| GPU GI CI PID Type Process name GPU Memory |\n",
+ "| ID ID Usage |\n",
+ "|=============================================================================|\n",
+ "| 0 N/A N/A 1232 C+G Insufficient Permissions N/A |\n",
+ "| 0 N/A N/A 1240 C+G Insufficient Permissions N/A |\n",
+ "| 0 N/A N/A 1528 C+G ...y\\ShellExperienceHost.exe N/A |\n",
+ "| 0 N/A N/A 7296 C+G Insufficient Permissions N/A |\n",
+ "| 0 N/A N/A 8280 C+G C:\\Windows\\explorer.exe N/A |\n",
+ "| 0 N/A N/A 9532 C+G ...artMenuExperienceHost.exe N/A |\n",
+ "| 0 N/A N/A 9896 C+G ...5n1h2txyewy\\SearchApp.exe N/A |\n",
+ "| 0 N/A N/A 11040 C+G ...2txyewy\\TextInputHost.exe N/A |\n",
+ "| 0 N/A N/A 11424 C+G Insufficient Permissions N/A |\n",
+ "| 0 N/A N/A 13112 C+G ...icrosoft VS Code\\Code.exe N/A |\n",
+ "| 0 N/A N/A 18720 C+G ...-2.9.15\\GitHubDesktop.exe N/A |\n",
+ "| 0 N/A N/A 22996 C+G ...bbwe\\Microsoft.Photos.exe N/A |\n",
+ "| 0 N/A N/A 23512 C+G ...me\\Application\\chrome.exe N/A |\n",
+ "| 0 N/A N/A 25892 C+G Insufficient Permissions N/A |\n",
+ "+-----------------------------------------------------------------------------+\n"
+ ]
+ }
+ ],
+ "source": [
+ "!nvidia-smi"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Z6BtQIgWuoqt"
+ },
+ "source": [
+ "Installation\n",
+ "All file changes made by this notebook are temporary. You can try to mount your own google drive to store files if you want."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "wdQJ9d8N8Tnf"
+ },
+ "source": [
+ "#Get Scripts"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "9jZWwt97uvIe"
+ },
+ "outputs": [],
+ "source": [
+ "!git clone https://github.com/neuralchen/SimSwap\n",
+ "!cd SimSwap && git pull"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ATLrrbso8Y-Y"
+ },
+ "source": [
+ "# Install Blocks"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "rwvbPhtOvZAL"
+ },
+ "outputs": [],
+ "source": [
+ "!pip install googledrivedownloader\n",
+ "!pip install timm==0.5.4\n",
+ "!wget -P SimSwap/arcface_model https://github.com/neuralchen/SimSwap/releases/download/1.0/arcface_checkpoint.tar"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "hleVtHIJ_QUK"
+ },
+ "source": [
+ "#Download the Training Dataset\n",
+ "We employ the cropped VGGFace2-224 dataset for this toy training demo.\n",
+ "\n",
+ "You can download the dataset from our google driver https://drive.google.com/file/d/19pWvdEHS-CEG6tW3PdxdtZ5QEymVjImc/view?usp=sharing\n",
+ "\n",
+ "***Please check the dataset in dir /content/TrainingData***\n",
+ "\n",
+ "***If dataset already exists in /content/TrainingData, please do not run blow scripts!***\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "h2tyjBl0Llxp"
+ },
+ "outputs": [],
+ "source": [
+ "%mkdir /content/TrainingData\n",
+ "!wget --load-cookies /tmp/cookies.txt \"https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=19pWvdEHS-CEG6tW3PdxdtZ5QEymVjImc' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p')&id=19pWvdEHS-CEG6tW3PdxdtZ5QEymVjImc\" -O /content/TrainingData/vggface2_crop_arcfacealign_224.tar && rm -rf /tmp/cookies.txt\n",
+ "%cd /content/\n",
+ "!tar -xzvf /content/TrainingData/vggface2_crop_arcfacealign_224.tar -C /content/TrainingData\n ",
+ "!rm /content/TrainingData/vggface2_crop_arcfacealign_224.tar"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "o5SNDWzA8LjJ"
+ },
+ "source": [
+ "#Trainig\n",
+ "Batch size must larger than 1!"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "XCxHa4oW507s"
+ },
+ "outputs": [],
+ "source": [
+ "%cd /content/SimSwap\n",
+ "!ls\n",
+ "!python train.py --name simswap224_test --gpu_ids 0 --dataset /content/TrainingData/vggface2_crop_arcfacealign_224 --Gdeep False"
+ ]
+ }
+ ],
+ "metadata": {
+ "accelerator": "GPU",
+ "colab": {
+ "collapsed_sections": [],
+ "name": "train.ipynb",
+ "provenance": []
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ },
+ "language_info": {
+ "name": "python",
+ "version": "3.8.5"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/SimSwap/train.py b/SimSwap/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..ebe5fdef594ed434b15ca8bfc909ac4e34d6d05b
--- /dev/null
+++ b/SimSwap/train.py
@@ -0,0 +1,290 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: train.py
+# Created Date: Monday December 27th 2021
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified: Friday, 22nd April 2022 10:49:26 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2021 Shanghai Jiao Tong University
+#############################################################
+
+import os
+import time
+import random
+import argparse
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from torch.backends import cudnn
+import torch.utils.tensorboard as tensorboard
+
+from util import util
+from util.plot import plot_batch
+
+from models.projected_model import fsModel
+from data.data_loader_Swapping import GetLoader
+
+def str2bool(v):
+ return v.lower() in ('true')
+
+class TrainOptions:
+ def __init__(self):
+ self.parser = argparse.ArgumentParser()
+ self.initialized = False
+
+ def initialize(self):
+ self.parser.add_argument('--name', type=str, default='simswap', help='name of the experiment. It decides where to store samples and models')
+ self.parser.add_argument('--gpu_ids', default='0')
+ self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
+ self.parser.add_argument('--isTrain', type=str2bool, default='True')
+
+ # input/output sizes
+ self.parser.add_argument('--batchSize', type=int, default=4, help='input batch size')
+
+ # for displays
+ self.parser.add_argument('--use_tensorboard', type=str2bool, default='False')
+
+ # for training
+ self.parser.add_argument('--dataset', type=str, default="/path/to/VGGFace2", help='path to the face swapping dataset')
+ self.parser.add_argument('--continue_train', type=str2bool, default='False', help='continue training: load the latest model')
+ self.parser.add_argument('--load_pretrain', type=str, default='./checkpoints/simswap224_test', help='load the pretrained model from the specified location')
+ self.parser.add_argument('--which_epoch', type=str, default='10000', help='which epoch to load? set to latest to use latest cached model')
+ self.parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
+ self.parser.add_argument('--niter', type=int, default=10000, help='# of iter at starting learning rate')
+ self.parser.add_argument('--niter_decay', type=int, default=10000, help='# of iter to linearly decay learning rate to zero')
+ self.parser.add_argument('--beta1', type=float, default=0.0, help='momentum term of adam')
+ self.parser.add_argument('--lr', type=float, default=0.0004, help='initial learning rate for adam')
+ self.parser.add_argument('--Gdeep', type=str2bool, default='False')
+
+ # for discriminators
+ self.parser.add_argument('--lambda_feat', type=float, default=10.0, help='weight for feature matching loss')
+ self.parser.add_argument('--lambda_id', type=float, default=30.0, help='weight for id loss')
+ self.parser.add_argument('--lambda_rec', type=float, default=10.0, help='weight for reconstruction loss')
+
+ self.parser.add_argument("--Arc_path", type=str, default='arcface_model/arcface_checkpoint.tar', help="run ONNX model via TRT")
+ self.parser.add_argument("--total_step", type=int, default=1000000, help='total training step')
+ self.parser.add_argument("--log_frep", type=int, default=200, help='frequence for printing log information')
+ self.parser.add_argument("--sample_freq", type=int, default=1000, help='frequence for sampling')
+ self.parser.add_argument("--model_freq", type=int, default=10000, help='frequence for saving the model')
+
+
+
+
+ self.isTrain = True
+
+ def parse(self, save=True):
+ if not self.initialized:
+ self.initialize()
+ self.opt = self.parser.parse_args()
+ self.opt.isTrain = self.isTrain # train or test
+
+ args = vars(self.opt)
+
+ print('------------ Options -------------')
+ for k, v in sorted(args.items()):
+ print('%s: %s' % (str(k), str(v)))
+ print('-------------- End ----------------')
+
+ # save to the disk
+ if self.opt.isTrain:
+ expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name)
+ util.mkdirs(expr_dir)
+ if save and not self.opt.continue_train:
+ file_name = os.path.join(expr_dir, 'opt.txt')
+ with open(file_name, 'wt') as opt_file:
+ opt_file.write('------------ Options -------------\n')
+ for k, v in sorted(args.items()):
+ opt_file.write('%s: %s\n' % (str(k), str(v)))
+ opt_file.write('-------------- End ----------------\n')
+ return self.opt
+
+
+if __name__ == '__main__':
+
+ opt = TrainOptions().parse()
+ iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt')
+
+ sample_path = os.path.join(opt.checkpoints_dir, opt.name, 'samples')
+
+ if not os.path.exists(sample_path):
+ os.makedirs(sample_path)
+
+ log_path = os.path.join(opt.checkpoints_dir, opt.name, 'summary')
+
+ if not os.path.exists(log_path):
+ os.makedirs(log_path)
+
+ if opt.continue_train:
+ try:
+ start_epoch, epoch_iter = np.loadtxt(iter_path , delimiter=',', dtype=int)
+ except:
+ start_epoch, epoch_iter = 1, 0
+ print('Resuming from epoch %d at iteration %d' % (start_epoch, epoch_iter))
+ else:
+ start_epoch, epoch_iter = 1, 0
+
+ os.environ['CUDA_VISIBLE_DEVICES'] = str(opt.gpu_ids)
+ print("GPU used : ", str(opt.gpu_ids))
+
+
+ cudnn.benchmark = True
+
+
+
+ model = fsModel()
+
+ model.initialize(opt)
+
+ #####################################################
+ if opt.use_tensorboard:
+ tensorboard_writer = tensorboard.SummaryWriter(log_path)
+ logger = tensorboard_writer
+
+ log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
+
+ with open(log_name, "a") as log_file:
+ now = time.strftime("%c")
+ log_file.write('================ Training Loss (%s) ================\n' % now)
+
+ optimizer_G, optimizer_D = model.optimizer_G, model.optimizer_D
+
+ loss_avg = 0
+ refresh_count = 0
+ imagenet_std = torch.Tensor([0.229, 0.224, 0.225]).view(3,1,1)
+ imagenet_mean = torch.Tensor([0.485, 0.456, 0.406]).view(3,1,1)
+
+ train_loader = GetLoader(opt.dataset,opt.batchSize,8,1234)
+
+ randindex = [i for i in range(opt.batchSize)]
+ random.shuffle(randindex)
+
+ if not opt.continue_train:
+ start = 0
+ else:
+ start = int(opt.which_epoch)
+ total_step = opt.total_step
+ import datetime
+ print("Start to train at %s"%(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')))
+
+ from util.logo_class import logo_class
+ logo_class.print_start_training()
+ model.netD.feature_network.requires_grad_(False)
+
+ # Training Cycle
+ for step in range(start, total_step):
+ model.netG.train()
+ for interval in range(2):
+ random.shuffle(randindex)
+ src_image1, src_image2 = train_loader.next()
+
+ if step%2 == 0:
+ img_id = src_image2
+ else:
+ img_id = src_image2[randindex]
+
+ img_id_112 = F.interpolate(img_id,size=(112,112), mode='bicubic')
+ latent_id = model.netArc(img_id_112)
+ latent_id = F.normalize(latent_id, p=2, dim=1)
+ if interval:
+
+ img_fake = model.netG(src_image1, latent_id)
+ gen_logits,_ = model.netD(img_fake.detach(), None)
+ loss_Dgen = (F.relu(torch.ones_like(gen_logits) + gen_logits)).mean()
+
+ real_logits,_ = model.netD(src_image2,None)
+ loss_Dreal = (F.relu(torch.ones_like(real_logits) - real_logits)).mean()
+
+ loss_D = loss_Dgen + loss_Dreal
+ optimizer_D.zero_grad()
+ loss_D.backward()
+ optimizer_D.step()
+ else:
+
+ # model.netD.requires_grad_(True)
+ img_fake = model.netG(src_image1, latent_id)
+ # G loss
+ gen_logits,feat = model.netD(img_fake, None)
+
+ loss_Gmain = (-gen_logits).mean()
+ img_fake_down = F.interpolate(img_fake, size=(112,112), mode='bicubic')
+ latent_fake = model.netArc(img_fake_down)
+ latent_fake = F.normalize(latent_fake, p=2, dim=1)
+ loss_G_ID = (1 - model.cosin_metric(latent_fake, latent_id)).mean()
+ real_feat = model.netD.get_feature(src_image1)
+ feat_match_loss = model.criterionFeat(feat["3"],real_feat["3"])
+ loss_G = loss_Gmain + loss_G_ID * opt.lambda_id + feat_match_loss * opt.lambda_feat
+
+
+ if step%2 == 0:
+ #G_Rec
+ loss_G_Rec = model.criterionRec(img_fake, src_image1) * opt.lambda_rec
+ loss_G += loss_G_Rec
+
+ optimizer_G.zero_grad()
+ loss_G.backward()
+ optimizer_G.step()
+
+
+ ############## Display results and errors ##########
+ ### print out errors
+ # Print out log info
+ if (step + 1) % opt.log_frep == 0:
+ # errors = {k: v.data.item() if not isinstance(v, int) else v for k, v in loss_dict.items()}
+ errors = {
+ "G_Loss":loss_Gmain.item(),
+ "G_ID":loss_G_ID.item(),
+ "G_Rec":loss_G_Rec.item(),
+ "G_feat_match":feat_match_loss.item(),
+ "D_fake":loss_Dgen.item(),
+ "D_real":loss_Dreal.item(),
+ "D_loss":loss_D.item()
+ }
+ if opt.use_tensorboard:
+ for tag, value in errors.items():
+ logger.add_scalar(tag, value, step)
+ message = '( step: %d, ) ' % (step)
+ for k, v in errors.items():
+ message += '%s: %.3f ' % (k, v)
+
+ print(message)
+ with open(log_name, "a") as log_file:
+ log_file.write('%s\n' % message)
+
+ ### display output images
+ if (step + 1) % opt.sample_freq == 0:
+ model.netG.eval()
+ with torch.no_grad():
+ imgs = list()
+ zero_img = (torch.zeros_like(src_image1[0,...]))
+ imgs.append(zero_img.cpu().numpy())
+ save_img = ((src_image1.cpu())* imagenet_std + imagenet_mean).numpy()
+ for r in range(opt.batchSize):
+ imgs.append(save_img[r,...])
+ arcface_112 = F.interpolate(src_image2,size=(112,112), mode='bicubic')
+ id_vector_src1 = model.netArc(arcface_112)
+ id_vector_src1 = F.normalize(id_vector_src1, p=2, dim=1)
+
+ for i in range(opt.batchSize):
+
+ imgs.append(save_img[i,...])
+ image_infer = src_image1[i, ...].repeat(opt.batchSize, 1, 1, 1)
+ img_fake = model.netG(image_infer, id_vector_src1).cpu()
+
+ img_fake = img_fake * imagenet_std
+ img_fake = img_fake + imagenet_mean
+ img_fake = img_fake.numpy()
+ for j in range(opt.batchSize):
+ imgs.append(img_fake[j,...])
+ print("Save test data")
+ imgs = np.stack(imgs, axis = 0).transpose(0,2,3,1)
+ plot_batch(imgs, os.path.join(sample_path, 'step_'+str(step+1)+'.jpg'))
+
+ ### save latest model
+ if (step+1) % opt.model_freq==0:
+ print('saving the latest model (steps %d)' % (step+1))
+ model.save(step+1)
+ np.savetxt(iter_path, (step+1, total_step), delimiter=',', fmt='%d')
+ wandb.finish()
\ No newline at end of file
diff --git a/SimSwap/util/__pycache__/add_watermark.cpython-310.pyc b/SimSwap/util/__pycache__/add_watermark.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..98c801fac408897a7f24f8c0b844fac8501557bc
Binary files /dev/null and b/SimSwap/util/__pycache__/add_watermark.cpython-310.pyc differ
diff --git a/SimSwap/util/__pycache__/norm.cpython-310.pyc b/SimSwap/util/__pycache__/norm.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7bd7ab017feda1f32d0f48b3c2ebc85687301ef9
Binary files /dev/null and b/SimSwap/util/__pycache__/norm.cpython-310.pyc differ
diff --git a/SimSwap/util/__pycache__/reverse2original.cpython-310.pyc b/SimSwap/util/__pycache__/reverse2original.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..af7ccc41d36b1b98a551bb11c66109ced82b5970
Binary files /dev/null and b/SimSwap/util/__pycache__/reverse2original.cpython-310.pyc differ
diff --git a/SimSwap/util/__pycache__/util.cpython-310.pyc b/SimSwap/util/__pycache__/util.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0b1ac559adaaa6986010901d4646e991f8a7147f
Binary files /dev/null and b/SimSwap/util/__pycache__/util.cpython-310.pyc differ
diff --git a/SimSwap/util/__pycache__/videoswap.cpython-310.pyc b/SimSwap/util/__pycache__/videoswap.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0e3a953941f453aad5defd4a14fb607b33675bd0
Binary files /dev/null and b/SimSwap/util/__pycache__/videoswap.cpython-310.pyc differ
diff --git a/SimSwap/util/add_watermark.py b/SimSwap/util/add_watermark.py
new file mode 100644
index 0000000000000000000000000000000000000000..a4c5d3d7f4d494fb12a3e2b32ca910431fdf61da
--- /dev/null
+++ b/SimSwap/util/add_watermark.py
@@ -0,0 +1,118 @@
+
+import cv2
+import numpy as np
+from PIL import Image
+import math
+import numpy as np
+# import torch
+# from torchvision import transforms
+
+def rotate_image(image, angle, center = None, scale = 1.0):
+ (h, w) = image.shape[:2]
+
+ if center is None:
+ center = (w / 2, h / 2)
+
+ # Perform the rotation
+ M = cv2.getRotationMatrix2D(center, angle, scale)
+ rotated = cv2.warpAffine(image, M, (w, h))
+
+ return rotated
+
+class watermark_image:
+ def __init__(self, logo_path, size=0.3, oritation="DR", margin=(5,20,20,20), angle=15, rgb_weight=(0,1,1.5), input_frame_shape=None) -> None:
+
+ logo_image = cv2.imread(logo_path, cv2.IMREAD_UNCHANGED)
+ h,w,c = logo_image.shape
+ if angle%360 != 0:
+ new_h = w*math.sin(angle/180*math.pi) + h*math.cos(angle/180*math.pi)
+ pad_h = int((new_h-h)//2)
+
+ padding = np.zeros((pad_h, w, c), dtype=np.uint8)
+ logo_image = cv2.vconcat([logo_image, padding])
+ logo_image = cv2.vconcat([padding, logo_image])
+
+ logo_image = rotate_image(logo_image, angle)
+ print(logo_image.shape)
+ self.logo_image = logo_image
+
+ if self.logo_image.shape[2] < 4:
+ print("No alpha channel found!")
+ self.logo_image = self.__addAlpha__(self.logo_image) #add alpha channel
+ self.size = size
+ self.oritation = oritation
+ self.margin = margin
+ self.ori_shape = self.logo_image.shape
+ self.resized = False
+ self.rgb_weight = rgb_weight
+
+ self.logo_image[:, :, 2] = self.logo_image[:, :, 2]*self.rgb_weight[0]
+ self.logo_image[:, :, 1] = self.logo_image[:, :, 1]*self.rgb_weight[1]
+ self.logo_image[:, :, 0] = self.logo_image[:, :, 0]*self.rgb_weight[2]
+
+ if input_frame_shape is not None:
+
+ logo_w = input_frame_shape[1] * self.size
+ ratio = logo_w / self.ori_shape[1]
+ logo_h = int(ratio * self.ori_shape[0])
+ logo_w = int(logo_w)
+
+ size = (logo_w, logo_h)
+ self.logo_image = cv2.resize(self.logo_image, size, interpolation = cv2.INTER_CUBIC)
+ self.resized = True
+ if oritation == "UL":
+ self.coor_h = self.margin[1]
+ self.coor_w = self.margin[0]
+ elif oritation == "UR":
+ self.coor_h = self.margin[1]
+ self.coor_w = input_frame_shape[1] - (logo_w + self.margin[2])
+ elif oritation == "DL":
+ self.coor_h = input_frame_shape[0] - (logo_h + self.margin[1])
+ self.coor_w = self.margin[0]
+ else:
+ self.coor_h = input_frame_shape[0] - (logo_h + self.margin[3])
+ self.coor_w = input_frame_shape[1] - (logo_w + self.margin[2])
+ self.logo_w = logo_w
+ self.logo_h = logo_h
+ self.mask = self.logo_image[:,:,3]
+ self.mask = cv2.bitwise_not(self.mask//255)
+
+ def apply_frames(self, frame):
+
+ if not self.resized:
+ shape = frame.shape
+ logo_w = shape[1] * self.size
+ ratio = logo_w / self.ori_shape[1]
+ logo_h = int(ratio * self.ori_shape[0])
+ logo_w = int(logo_w)
+
+ size = (logo_w, logo_h)
+ self.logo_image = cv2.resize(self.logo_image, size, interpolation = cv2.INTER_CUBIC)
+ self.resized = True
+ if self.oritation == "UL":
+ self.coor_h = self.margin[1]
+ self.coor_w = self.margin[0]
+ elif self.oritation == "UR":
+ self.coor_h = self.margin[1]
+ self.coor_w = shape[1] - (logo_w + self.margin[2])
+ elif self.oritation == "DL":
+ self.coor_h = shape[0] - (logo_h + self.margin[1])
+ self.coor_w = self.margin[0]
+ else:
+ self.coor_h = shape[0] - (logo_h + self.margin[3])
+ self.coor_w = shape[1] - (logo_w + self.margin[2])
+ self.logo_w = logo_w
+ self.logo_h = logo_h
+ self.mask = self.logo_image[:,:,3]
+ self.mask = cv2.bitwise_not(self.mask//255)
+
+ original_frame = frame[self.coor_h:(self.coor_h+self.logo_h), self.coor_w:(self.coor_w+self.logo_w),:]
+ blending_logo = cv2.add(self.logo_image[:,:,0:3],original_frame,mask = self.mask)
+ frame[self.coor_h:(self.coor_h+self.logo_h), self.coor_w:(self.coor_w+self.logo_w),:] = blending_logo
+ return frame
+
+ def __addAlpha__(self, image):
+ shape = image.shape
+ alpha_channel = np.ones((shape[0],shape[1],1),np.uint8)*255
+ return np.concatenate((image,alpha_channel),2)
+
diff --git a/SimSwap/util/html.py b/SimSwap/util/html.py
new file mode 100644
index 0000000000000000000000000000000000000000..71c48ad4070da60f2ed72217432c34c7f5e5fef0
--- /dev/null
+++ b/SimSwap/util/html.py
@@ -0,0 +1,63 @@
+import dominate
+from dominate.tags import *
+import os
+
+
+class HTML:
+ def __init__(self, web_dir, title, refresh=0):
+ self.title = title
+ self.web_dir = web_dir
+ self.img_dir = os.path.join(self.web_dir, 'images')
+ if not os.path.exists(self.web_dir):
+ os.makedirs(self.web_dir)
+ if not os.path.exists(self.img_dir):
+ os.makedirs(self.img_dir)
+
+ self.doc = dominate.document(title=title)
+ if refresh > 0:
+ with self.doc.head:
+ meta(http_equiv="refresh", content=str(refresh))
+
+ def get_image_dir(self):
+ return self.img_dir
+
+ def add_header(self, str):
+ with self.doc:
+ h3(str)
+
+ def add_table(self, border=1):
+ self.t = table(border=border, style="table-layout: fixed;")
+ self.doc.add(self.t)
+
+ def add_images(self, ims, txts, links, width=512):
+ self.add_table()
+ with self.t:
+ with tr():
+ for im, txt, link in zip(ims, txts, links):
+ with td(style="word-wrap: break-word;", halign="center", valign="top"):
+ with p():
+ with a(href=os.path.join('images', link)):
+ img(style="width:%dpx" % (width), src=os.path.join('images', im))
+ br()
+ p(txt)
+
+ def save(self):
+ html_file = '%s/index.html' % self.web_dir
+ f = open(html_file, 'wt')
+ f.write(self.doc.render())
+ f.close()
+
+
+if __name__ == '__main__':
+ html = HTML('web/', 'test_html')
+ html.add_header('hello world')
+
+ ims = []
+ txts = []
+ links = []
+ for n in range(4):
+ ims.append('image_%d.jpg' % n)
+ txts.append('text_%d' % n)
+ links.append('image_%d.jpg' % n)
+ html.add_images(ims, txts, links)
+ html.save()
diff --git a/SimSwap/util/image_pool.py b/SimSwap/util/image_pool.py
new file mode 100644
index 0000000000000000000000000000000000000000..63e18778da7b9e029a8d536a3c4329c33e17b4d1
--- /dev/null
+++ b/SimSwap/util/image_pool.py
@@ -0,0 +1,31 @@
+import random
+import torch
+from torch.autograd import Variable
+class ImagePool():
+ def __init__(self, pool_size):
+ self.pool_size = pool_size
+ if self.pool_size > 0:
+ self.num_imgs = 0
+ self.images = []
+
+ def query(self, images):
+ if self.pool_size == 0:
+ return images
+ return_images = []
+ for image in images.data:
+ image = torch.unsqueeze(image, 0)
+ if self.num_imgs < self.pool_size:
+ self.num_imgs = self.num_imgs + 1
+ self.images.append(image)
+ return_images.append(image)
+ else:
+ p = random.uniform(0, 1)
+ if p > 0.5:
+ random_id = random.randint(0, self.pool_size-1)
+ tmp = self.images[random_id].clone()
+ self.images[random_id] = image
+ return_images.append(tmp)
+ else:
+ return_images.append(image)
+ return_images = Variable(torch.cat(return_images, 0))
+ return return_images
diff --git a/SimSwap/util/json_config.py b/SimSwap/util/json_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..c68fbffd667852487084106527605896e6dfa0d6
--- /dev/null
+++ b/SimSwap/util/json_config.py
@@ -0,0 +1,15 @@
+import json
+
+
+def readConfig(path):
+ with open(path,'r') as cf:
+ nodelocaltionstr = cf.read()
+ nodelocaltioninf = json.loads(nodelocaltionstr)
+ if isinstance(nodelocaltioninf,str):
+ nodelocaltioninf = json.loads(nodelocaltioninf)
+ return nodelocaltioninf
+
+def writeConfig(path, info):
+ with open(path, 'w') as cf:
+ configjson = json.dumps(info, indent=4)
+ cf.writelines(configjson)
\ No newline at end of file
diff --git a/SimSwap/util/logo_class.py b/SimSwap/util/logo_class.py
new file mode 100644
index 0000000000000000000000000000000000000000..044dce3c610c8b03a3313f65bd2a87188a092888
--- /dev/null
+++ b/SimSwap/util/logo_class.py
@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: logo_class.py
+# Created Date: Tuesday June 29th 2021
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified: Monday, 11th October 2021 12:39:55 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2021 Shanghai Jiao Tong University
+#############################################################
+
+class logo_class:
+
+ @staticmethod
+ def print_group_logo():
+ logo_str = """
+
+███╗ ██╗██████╗ ███████╗██╗ ██████╗ ███████╗ ██╗████████╗██╗ ██╗
+████╗ ██║██╔══██╗██╔════╝██║██╔════╝ ██╔════╝ ██║╚══██╔══╝██║ ██║
+██╔██╗ ██║██████╔╝███████╗██║██║ ███╗ ███████╗ ██║ ██║ ██║ ██║
+██║╚██╗██║██╔══██╗╚════██║██║██║ ██║ ╚════██║██ ██║ ██║ ██║ ██║
+██║ ╚████║██║ ██║███████║██║╚██████╔╝ ███████║╚█████╔╝ ██║ ╚██████╔╝
+╚═╝ ╚═══╝╚═╝ ╚═╝╚══════╝╚═╝ ╚═════╝ ╚══════╝ ╚════╝ ╚═╝ ╚═════╝
+Neural Rendering Special Interesting Group of SJTU
+
+ """
+ print(logo_str)
+
+ @staticmethod
+ def print_start_training():
+ logo_str = """
+ _____ __ __ ______ _ _
+ / ___/ / /_ ____ _ _____ / /_ /_ __/_____ ____ _ (_)____ (_)____ ____ _
+ \__ \ / __// __ `// ___// __/ / / / ___// __ `// // __ \ / // __ \ / __ `/
+ ___/ // /_ / /_/ // / / /_ / / / / / /_/ // // / / // // / / // /_/ /
+/____/ \__/ \__,_//_/ \__/ /_/ /_/ \__,_//_//_/ /_//_//_/ /_/ \__, /
+ /____/
+ """
+ print(logo_str)
+
+if __name__=="__main__":
+ # logo_class.print_group_logo()
+ logo_class.print_start_training()
\ No newline at end of file
diff --git a/SimSwap/util/norm.py b/SimSwap/util/norm.py
new file mode 100644
index 0000000000000000000000000000000000000000..981adcf2654c353ae63ac2b6cca5d70aa14eb46f
--- /dev/null
+++ b/SimSwap/util/norm.py
@@ -0,0 +1,25 @@
+import torch.nn as nn
+import numpy as np
+import torch
+class SpecificNorm(nn.Module):
+ def __init__(self, epsilon=1e-8):
+ """
+ @notice: avoid in-place ops.
+ https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
+ """
+ super(SpecificNorm, self).__init__()
+ self.mean = np.array([0.485, 0.456, 0.406])
+ self.mean = torch.from_numpy(self.mean).float().cuda()
+ self.mean = self.mean.view([1, 3, 1, 1])
+
+ self.std = np.array([0.229, 0.224, 0.225])
+ self.std = torch.from_numpy(self.std).float().cuda()
+ self.std = self.std.view([1, 3, 1, 1])
+
+ def forward(self, x):
+ mean = self.mean.expand([1, 3, x.shape[2], x.shape[3]])
+ std = self.std.expand([1, 3, x.shape[2], x.shape[3]])
+
+ x = (x - mean) / std
+
+ return x
\ No newline at end of file
diff --git a/SimSwap/util/plot.py b/SimSwap/util/plot.py
new file mode 100644
index 0000000000000000000000000000000000000000..0da1c75b8dfd4c073d7c4c1686e76b8d0bab51af
--- /dev/null
+++ b/SimSwap/util/plot.py
@@ -0,0 +1,37 @@
+import numpy as np
+import math
+import PIL
+
+def postprocess(x):
+ """[0,1] to uint8."""
+
+ x = np.clip(255 * x, 0, 255)
+ x = np.cast[np.uint8](x)
+ return x
+
+def tile(X, rows, cols):
+ """Tile images for display."""
+ tiling = np.zeros((rows * X.shape[1], cols * X.shape[2], X.shape[3]), dtype = X.dtype)
+ for i in range(rows):
+ for j in range(cols):
+ idx = i * cols + j
+ if idx < X.shape[0]:
+ img = X[idx,...]
+ tiling[
+ i*X.shape[1]:(i+1)*X.shape[1],
+ j*X.shape[2]:(j+1)*X.shape[2],
+ :] = img
+ return tiling
+
+
+def plot_batch(X, out_path):
+ """Save batch of images tiled."""
+ n_channels = X.shape[3]
+ if n_channels > 3:
+ X = X[:,:,:,np.random.choice(n_channels, size = 3)]
+ X = postprocess(X)
+ rc = math.sqrt(X.shape[0])
+ rows = cols = math.ceil(rc)
+ canvas = tile(X, rows, cols)
+ canvas = np.squeeze(canvas)
+ PIL.Image.fromarray(canvas).save(out_path)
\ No newline at end of file
diff --git a/SimSwap/util/reverse2original.py b/SimSwap/util/reverse2original.py
new file mode 100644
index 0000000000000000000000000000000000000000..a74c998f706111cc46975091dd5dde52e43320af
--- /dev/null
+++ b/SimSwap/util/reverse2original.py
@@ -0,0 +1,177 @@
+import cv2
+import numpy as np
+# import time
+import torch
+from torch.nn import functional as F
+import torch.nn as nn
+
+
+def encode_segmentation_rgb(segmentation, no_neck=True):
+ parse = segmentation
+
+ face_part_ids = [1, 2, 3, 4, 5, 6, 10, 12, 13] if no_neck else [1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 13, 14]
+ mouth_id = 11
+ # hair_id = 17
+ face_map = np.zeros([parse.shape[0], parse.shape[1]])
+ mouth_map = np.zeros([parse.shape[0], parse.shape[1]])
+ # hair_map = np.zeros([parse.shape[0], parse.shape[1]])
+
+ for valid_id in face_part_ids:
+ valid_index = np.where(parse==valid_id)
+ face_map[valid_index] = 255
+ valid_index = np.where(parse==mouth_id)
+ mouth_map[valid_index] = 255
+ # valid_index = np.where(parse==hair_id)
+ # hair_map[valid_index] = 255
+ #return np.stack([face_map, mouth_map,hair_map], axis=2)
+ return np.stack([face_map, mouth_map], axis=2)
+
+
+class SoftErosion(nn.Module):
+ def __init__(self, kernel_size=15, threshold=0.6, iterations=1):
+ super(SoftErosion, self).__init__()
+ r = kernel_size // 2
+ self.padding = r
+ self.iterations = iterations
+ self.threshold = threshold
+
+ # Create kernel
+ y_indices, x_indices = torch.meshgrid(torch.arange(0., kernel_size), torch.arange(0., kernel_size))
+ dist = torch.sqrt((x_indices - r) ** 2 + (y_indices - r) ** 2)
+ kernel = dist.max() - dist
+ kernel /= kernel.sum()
+ kernel = kernel.view(1, 1, *kernel.shape)
+ self.register_buffer('weight', kernel)
+
+ def forward(self, x):
+ x = x.float()
+ for i in range(self.iterations - 1):
+ x = torch.min(x, F.conv2d(x, weight=self.weight, groups=x.shape[1], padding=self.padding))
+ x = F.conv2d(x, weight=self.weight, groups=x.shape[1], padding=self.padding)
+
+ mask = x >= self.threshold
+ x[mask] = 1.0
+ x[~mask] /= x[~mask].max()
+
+ return x, mask
+
+
+def postprocess(swapped_face, target, target_mask,smooth_mask):
+ # target_mask = cv2.resize(target_mask, (self.size, self.size))
+
+ mask_tensor = torch.from_numpy(target_mask.copy().transpose((2, 0, 1))).float().mul_(1/255.0).cuda()
+ face_mask_tensor = mask_tensor[0] + mask_tensor[1]
+
+ soft_face_mask_tensor, _ = smooth_mask(face_mask_tensor.unsqueeze_(0).unsqueeze_(0))
+ soft_face_mask_tensor.squeeze_()
+
+ soft_face_mask = soft_face_mask_tensor.cpu().numpy()
+ soft_face_mask = soft_face_mask[:, :, np.newaxis]
+
+ result = swapped_face * soft_face_mask + target * (1 - soft_face_mask)
+ result = result[:,:,::-1]# .astype(np.uint8)
+ return result
+
+def reverse2wholeimage(b_align_crop_tenor_list,swaped_imgs, mats, crop_size, oriimg, logoclass, save_path = '', \
+ no_simswaplogo = False,pasring_model =None,norm = None, use_mask = False):
+
+ target_image_list = []
+ img_mask_list = []
+ if use_mask:
+ smooth_mask = SoftErosion(kernel_size=17, threshold=0.9, iterations=7).cuda()
+ else:
+ pass
+
+ # print(len(swaped_imgs))
+ # print(mats)
+ # print(len(b_align_crop_tenor_list))
+ for swaped_img, mat ,source_img in zip(swaped_imgs, mats,b_align_crop_tenor_list):
+ swaped_img = swaped_img.cpu().detach().numpy().transpose((1, 2, 0))
+ img_white = np.full((crop_size,crop_size), 255, dtype=float)
+
+ # inverse the Affine transformation matrix
+ mat_rev = np.zeros([2,3])
+ div1 = mat[0][0]*mat[1][1]-mat[0][1]*mat[1][0]
+ mat_rev[0][0] = mat[1][1]/div1
+ mat_rev[0][1] = -mat[0][1]/div1
+ mat_rev[0][2] = -(mat[0][2]*mat[1][1]-mat[0][1]*mat[1][2])/div1
+ div2 = mat[0][1]*mat[1][0]-mat[0][0]*mat[1][1]
+ mat_rev[1][0] = mat[1][0]/div2
+ mat_rev[1][1] = -mat[0][0]/div2
+ mat_rev[1][2] = -(mat[0][2]*mat[1][0]-mat[0][0]*mat[1][2])/div2
+
+ orisize = (oriimg.shape[1], oriimg.shape[0])
+ if use_mask:
+ source_img_norm = norm(source_img)
+ source_img_512 = F.interpolate(source_img_norm,size=(512,512))
+ out = pasring_model(source_img_512)[0]
+ parsing = out.squeeze(0).detach().cpu().numpy().argmax(0)
+ vis_parsing_anno = parsing.copy().astype(np.uint8)
+ tgt_mask = encode_segmentation_rgb(vis_parsing_anno)
+ if tgt_mask.sum() >= 5000:
+ # face_mask_tensor = tgt_mask[...,0] + tgt_mask[...,1]
+ target_mask = cv2.resize(tgt_mask, (crop_size, crop_size))
+ # print(source_img)
+ target_image_parsing = postprocess(swaped_img, source_img[0].cpu().detach().numpy().transpose((1, 2, 0)), target_mask,smooth_mask)
+
+
+ target_image = cv2.warpAffine(target_image_parsing, mat_rev, orisize)
+ # target_image_parsing = cv2.warpAffine(swaped_img, mat_rev, orisize)
+ else:
+ target_image = cv2.warpAffine(swaped_img, mat_rev, orisize)[..., ::-1]
+ else:
+ target_image = cv2.warpAffine(swaped_img, mat_rev, orisize)
+ # source_image = cv2.warpAffine(source_img, mat_rev, orisize)
+
+ img_white = cv2.warpAffine(img_white, mat_rev, orisize)
+
+
+ img_white[img_white>20] =255
+
+ img_mask = img_white
+
+ # if use_mask:
+ # kernel = np.ones((40,40),np.uint8)
+ # img_mask = cv2.erode(img_mask,kernel,iterations = 1)
+ # else:
+ kernel = np.ones((40,40),np.uint8)
+ img_mask = cv2.erode(img_mask,kernel,iterations = 1)
+ kernel_size = (20, 20)
+ blur_size = tuple(2*i+1 for i in kernel_size)
+ img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
+
+ # kernel = np.ones((10,10),np.uint8)
+ # img_mask = cv2.erode(img_mask,kernel,iterations = 1)
+
+
+
+ img_mask /= 255
+
+ img_mask = np.reshape(img_mask, [img_mask.shape[0],img_mask.shape[1],1])
+
+ # pasing mask
+
+ # target_image_parsing = postprocess(target_image, source_image, tgt_mask)
+
+ if use_mask:
+ # target_image = np.array(target_image, dtype=np.float) * 255
+ target_image = np.array(target_image, dtype=float) * 255
+
+ else:
+ target_image = np.array(target_image, dtype=float)[..., ::-1] * 255
+
+
+ img_mask_list.append(img_mask)
+ target_image_list.append(target_image)
+
+
+ # target_image /= 255
+ # target_image = 0
+ img = np.array(oriimg, dtype=float)
+ for img_mask, target_image in zip(img_mask_list, target_image_list):
+ img = img_mask * target_image + (1-img_mask) * img
+
+ final_img = img.astype(np.uint8)
+ if not no_simswaplogo:
+ final_img = logoclass.apply_frames(final_img)
+ cv2.imwrite(save_path, final_img)
diff --git a/SimSwap/util/save_heatmap.py b/SimSwap/util/save_heatmap.py
new file mode 100644
index 0000000000000000000000000000000000000000..71ce4c985df204394f4556b71205511f241cba5f
--- /dev/null
+++ b/SimSwap/util/save_heatmap.py
@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: save_heatmap.py
+# Created Date: Friday January 15th 2021
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified: Wednesday, 19th January 2022 1:22:47 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2021 Shanghai Jiao Tong University
+#############################################################
+
+import os
+import shutil
+import seaborn as sns
+import matplotlib.pyplot as plt
+import cv2
+import numpy as np
+
+def SaveHeatmap(heatmaps, path, row=-1, dpi=72):
+ """
+ The input tensor must be B X 1 X H X W
+ """
+ batch_size = heatmaps.shape[0]
+ temp_path = ".temp/"
+ if not os.path.exists(temp_path):
+ os.makedirs(temp_path)
+ final_img = None
+ if row < 1:
+ col = batch_size
+ row = 1
+ else:
+ col = batch_size // row
+ if row * col
= col:
+ col_i = 0
+ row_i += 1
+ cv2.imwrite(path,final_img)
+
+if __name__ == "__main__":
+ random_map = np.random.randn(16,1,10,10)
+ SaveHeatmap(random_map,"./wocao.png",1)
diff --git a/SimSwap/util/util.py b/SimSwap/util/util.py
new file mode 100644
index 0000000000000000000000000000000000000000..f4f79ec0dfc7e7d48a85cf6d3d2e48cffcf7c731
--- /dev/null
+++ b/SimSwap/util/util.py
@@ -0,0 +1,100 @@
+from __future__ import print_function
+import torch
+import numpy as np
+from PIL import Image
+import numpy as np
+import os
+
+# Converts a Tensor into a Numpy array
+# |imtype|: the desired type of the converted numpy array
+def tensor2im(image_tensor, imtype=np.uint8, normalize=True):
+ if isinstance(image_tensor, list):
+ image_numpy = []
+ for i in range(len(image_tensor)):
+ image_numpy.append(tensor2im(image_tensor[i], imtype, normalize))
+ return image_numpy
+ image_numpy = image_tensor.cpu().float().numpy()
+ if normalize:
+ image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
+ else:
+ image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0
+ image_numpy = np.clip(image_numpy, 0, 255)
+ if image_numpy.shape[2] == 1 or image_numpy.shape[2] > 3:
+ image_numpy = image_numpy[:,:,0]
+ return image_numpy.astype(imtype)
+
+# Converts a one-hot tensor into a colorful label map
+def tensor2label(label_tensor, n_label, imtype=np.uint8):
+ if n_label == 0:
+ return tensor2im(label_tensor, imtype)
+ label_tensor = label_tensor.cpu().float()
+ if label_tensor.size()[0] > 1:
+ label_tensor = label_tensor.max(0, keepdim=True)[1]
+ label_tensor = Colorize(n_label)(label_tensor)
+ label_numpy = np.transpose(label_tensor.numpy(), (1, 2, 0))
+ return label_numpy.astype(imtype)
+
+def save_image(image_numpy, image_path):
+ image_pil = Image.fromarray(image_numpy)
+ image_pil.save(image_path)
+
+def mkdirs(paths):
+ if isinstance(paths, list) and not isinstance(paths, str):
+ for path in paths:
+ mkdir(path)
+ else:
+ mkdir(paths)
+
+def mkdir(path):
+ if not os.path.exists(path):
+ os.makedirs(path)
+
+###############################################################################
+# Code from
+# https://github.com/ycszen/pytorch-seg/blob/master/transform.py
+# Modified so it complies with the Citscape label map colors
+###############################################################################
+def uint82bin(n, count=8):
+ """returns the binary of integer n, count refers to amount of bits"""
+ return ''.join([str((n >> y) & 1) for y in range(count-1, -1, -1)])
+
+def labelcolormap(N):
+ if N == 35: # cityscape
+ cmap = np.array([( 0, 0, 0), ( 0, 0, 0), ( 0, 0, 0), ( 0, 0, 0), ( 0, 0, 0), (111, 74, 0), ( 81, 0, 81),
+ (128, 64,128), (244, 35,232), (250,170,160), (230,150,140), ( 70, 70, 70), (102,102,156), (190,153,153),
+ (180,165,180), (150,100,100), (150,120, 90), (153,153,153), (153,153,153), (250,170, 30), (220,220, 0),
+ (107,142, 35), (152,251,152), ( 70,130,180), (220, 20, 60), (255, 0, 0), ( 0, 0,142), ( 0, 0, 70),
+ ( 0, 60,100), ( 0, 0, 90), ( 0, 0,110), ( 0, 80,100), ( 0, 0,230), (119, 11, 32), ( 0, 0,142)],
+ dtype=np.uint8)
+ else:
+ cmap = np.zeros((N, 3), dtype=np.uint8)
+ for i in range(N):
+ r, g, b = 0, 0, 0
+ id = i
+ for j in range(7):
+ str_id = uint82bin(id)
+ r = r ^ (np.uint8(str_id[-1]) << (7-j))
+ g = g ^ (np.uint8(str_id[-2]) << (7-j))
+ b = b ^ (np.uint8(str_id[-3]) << (7-j))
+ id = id >> 3
+ cmap[i, 0] = r
+ cmap[i, 1] = g
+ cmap[i, 2] = b
+ return cmap
+
+class Colorize(object):
+ def __init__(self, n=35):
+ self.cmap = labelcolormap(n)
+ self.cmap = torch.from_numpy(self.cmap[:n])
+
+ def __call__(self, gray_image):
+ size = gray_image.size()
+ color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0)
+
+ for label in range(0, len(self.cmap)):
+ mask = (label == gray_image[0]).cpu()
+ color_image[0][mask] = self.cmap[label][0]
+ color_image[1][mask] = self.cmap[label][1]
+ color_image[2][mask] = self.cmap[label][2]
+
+ return color_image
diff --git a/SimSwap/util/videoswap.py b/SimSwap/util/videoswap.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a6c8fe58980be05477ab07294cacdc5d871eb80
--- /dev/null
+++ b/SimSwap/util/videoswap.py
@@ -0,0 +1,122 @@
+'''
+Author: Naiyuan liu
+Github: https://github.com/NNNNAI
+Date: 2021-11-23 17:03:58
+LastEditors: Naiyuan liu
+LastEditTime: 2021-11-24 19:19:52
+Description:
+'''
+import os
+import cv2
+import glob
+import torch
+import shutil
+import numpy as np
+from tqdm import tqdm
+from util.reverse2original import reverse2wholeimage
+import moviepy.editor as mp
+from moviepy.editor import AudioFileClip, VideoFileClip
+from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
+import time
+from util.add_watermark import watermark_image
+from util.norm import SpecificNorm
+from parsing_model.model import BiSeNet
+
+def _totensor(array):
+ tensor = torch.from_numpy(array)
+ img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
+ return img.float().div(255)
+
+def video_swap(video_path, id_vetor, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224, no_simswaplogo = False,use_mask =False):
+ video_forcheck = VideoFileClip(video_path)
+ if video_forcheck.audio is None:
+ no_audio = True
+ else:
+ no_audio = False
+
+ del video_forcheck
+
+ if not no_audio:
+ video_audio_clip = AudioFileClip(video_path)
+
+ video = cv2.VideoCapture(video_path)
+ logoclass = watermark_image('./simswaplogo/simswaplogo.png')
+ ret = True
+ frame_index = 0
+
+ frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+
+ # video_WIDTH = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
+
+ # video_HEIGHT = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+ fps = video.get(cv2.CAP_PROP_FPS)
+ if os.path.exists(temp_results_dir):
+ shutil.rmtree(temp_results_dir)
+
+ spNorm =SpecificNorm()
+ if use_mask:
+ n_classes = 19
+ net = BiSeNet(n_classes=n_classes)
+ net.cuda()
+ save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+ net.load_state_dict(torch.load(save_pth))
+ net.eval()
+ else:
+ net =None
+
+ # while ret:
+ for frame_index in tqdm(range(frame_count)):
+ ret, frame = video.read()
+ if ret:
+ detect_results = detect_model.get(frame,crop_size)
+
+ if detect_results is not None:
+ # print(frame_index)
+ if not os.path.exists(temp_results_dir):
+ os.mkdir(temp_results_dir)
+ frame_align_crop_list = detect_results[0]
+ frame_mat_list = detect_results[1]
+ swap_result_list = []
+ frame_align_crop_tenor_list = []
+ for frame_align_crop in frame_align_crop_list:
+
+ # BGR TO RGB
+ # frame_align_crop_RGB = frame_align_crop[...,::-1]
+
+ frame_align_crop_tenor = _totensor(cv2.cvtColor(frame_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
+
+ swap_result = swap_model(None, frame_align_crop_tenor, id_vetor, None, True)[0]
+ cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame)
+ swap_result_list.append(swap_result)
+ frame_align_crop_tenor_list.append(frame_align_crop_tenor)
+
+
+
+ reverse2wholeimage(frame_align_crop_tenor_list,swap_result_list, frame_mat_list, crop_size, frame, logoclass,\
+ os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo,pasring_model =net,use_mask=use_mask, norm = spNorm)
+
+ else:
+ if not os.path.exists(temp_results_dir):
+ os.mkdir(temp_results_dir)
+ frame = frame.astype(np.uint8)
+ if not no_simswaplogo:
+ frame = logoclass.apply_frames(frame)
+ cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame)
+ else:
+ break
+
+ video.release()
+
+ # image_filename_list = []
+ path = os.path.join(temp_results_dir,'*.jpg')
+ image_filenames = sorted(glob.glob(path))
+
+ clips = ImageSequenceClip(image_filenames,fps = fps)
+
+ if not no_audio:
+ clips = clips.set_audio(video_audio_clip)
+
+
+ clips.write_videofile(save_path,audio_codec='aac')
+
diff --git a/SimSwap/util/videoswap_multispecific.py b/SimSwap/util/videoswap_multispecific.py
new file mode 100644
index 0000000000000000000000000000000000000000..40e94880c31076da5c7eab000be916d26f9b30fc
--- /dev/null
+++ b/SimSwap/util/videoswap_multispecific.py
@@ -0,0 +1,146 @@
+import os
+import cv2
+import glob
+import torch
+import shutil
+import numpy as np
+from tqdm import tqdm
+from util.reverse2original import reverse2wholeimage
+import moviepy.editor as mp
+from moviepy.editor import AudioFileClip, VideoFileClip
+from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
+import time
+from util.add_watermark import watermark_image
+from util.norm import SpecificNorm
+import torch.nn.functional as F
+from parsing_model.model import BiSeNet
+
+def _totensor(array):
+ tensor = torch.from_numpy(array)
+ img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
+ return img.float().div(255)
+
+def video_swap(video_path, target_id_norm_list,source_specific_id_nonorm_list,id_thres, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224, no_simswaplogo = False,use_mask =False):
+ video_forcheck = VideoFileClip(video_path)
+ if video_forcheck.audio is None:
+ no_audio = True
+ else:
+ no_audio = False
+
+ del video_forcheck
+
+ if not no_audio:
+ video_audio_clip = AudioFileClip(video_path)
+
+ video = cv2.VideoCapture(video_path)
+ logoclass = watermark_image('./simswaplogo/simswaplogo.png')
+ ret = True
+ frame_index = 0
+
+ frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+
+ # video_WIDTH = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
+
+ # video_HEIGHT = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+ fps = video.get(cv2.CAP_PROP_FPS)
+ if os.path.exists(temp_results_dir):
+ shutil.rmtree(temp_results_dir)
+
+ spNorm =SpecificNorm()
+ mse = torch.nn.MSELoss().cuda()
+
+ if use_mask:
+ n_classes = 19
+ net = BiSeNet(n_classes=n_classes)
+ net.cuda()
+ save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+ net.load_state_dict(torch.load(save_pth))
+ net.eval()
+ else:
+ net =None
+
+ # while ret:
+ for frame_index in tqdm(range(frame_count)):
+ ret, frame = video.read()
+ if ret:
+ detect_results = detect_model.get(frame,crop_size)
+
+ if detect_results is not None:
+ # print(frame_index)
+ if not os.path.exists(temp_results_dir):
+ os.mkdir(temp_results_dir)
+ frame_align_crop_list = detect_results[0]
+ frame_mat_list = detect_results[1]
+
+ id_compare_values = []
+ frame_align_crop_tenor_list = []
+ for frame_align_crop in frame_align_crop_list:
+
+ # BGR TO RGB
+ # frame_align_crop_RGB = frame_align_crop[...,::-1]
+
+ frame_align_crop_tenor = _totensor(cv2.cvtColor(frame_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
+
+ frame_align_crop_tenor_arcnorm = spNorm(frame_align_crop_tenor)
+ frame_align_crop_tenor_arcnorm_downsample = F.interpolate(frame_align_crop_tenor_arcnorm, size=(112,112))
+ frame_align_crop_crop_id_nonorm = swap_model.netArc(frame_align_crop_tenor_arcnorm_downsample)
+ id_compare_values.append([])
+ for source_specific_id_nonorm_tmp in source_specific_id_nonorm_list:
+ id_compare_values[-1].append(mse(frame_align_crop_crop_id_nonorm,source_specific_id_nonorm_tmp).detach().cpu().numpy())
+ frame_align_crop_tenor_list.append(frame_align_crop_tenor)
+
+ id_compare_values_array = np.array(id_compare_values).transpose(1,0)
+ min_indexs = np.argmin(id_compare_values_array,axis=0)
+ min_value = np.min(id_compare_values_array,axis=0)
+
+ swap_result_list = []
+ swap_result_matrix_list = []
+ swap_result_ori_pic_list = []
+ for tmp_index, min_index in enumerate(min_indexs):
+ if min_value[tmp_index] < id_thres:
+ swap_result = swap_model(None, frame_align_crop_tenor_list[tmp_index], target_id_norm_list[min_index], None, True)[0]
+ swap_result_list.append(swap_result)
+ swap_result_matrix_list.append(frame_mat_list[tmp_index])
+ swap_result_ori_pic_list.append(frame_align_crop_tenor_list[tmp_index])
+ else:
+ pass
+
+
+
+ if len(swap_result_list) !=0:
+
+ reverse2wholeimage(swap_result_ori_pic_list,swap_result_list, swap_result_matrix_list, crop_size, frame, logoclass,\
+ os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo,pasring_model =net,use_mask=use_mask, norm = spNorm)
+ else:
+ if not os.path.exists(temp_results_dir):
+ os.mkdir(temp_results_dir)
+ frame = frame.astype(np.uint8)
+ if not no_simswaplogo:
+ frame = logoclass.apply_frames(frame)
+ cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame)
+
+ else:
+ if not os.path.exists(temp_results_dir):
+ os.mkdir(temp_results_dir)
+ frame = frame.astype(np.uint8)
+ if not no_simswaplogo:
+ frame = logoclass.apply_frames(frame)
+ cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame)
+ else:
+ break
+
+ video.release()
+
+ # image_filename_list = []
+ path = os.path.join(temp_results_dir,'*.jpg')
+ image_filenames = sorted(glob.glob(path))
+
+ clips = ImageSequenceClip(image_filenames,fps = fps)
+
+ if not no_audio:
+ clips = clips.set_audio(video_audio_clip)
+
+
+ clips.write_videofile(save_path,audio_codec='aac')
+
diff --git a/SimSwap/util/videoswap_specific.py b/SimSwap/util/videoswap_specific.py
new file mode 100644
index 0000000000000000000000000000000000000000..8177efb171ec4ad1ab6fa66aa3ca47b166c11760
--- /dev/null
+++ b/SimSwap/util/videoswap_specific.py
@@ -0,0 +1,130 @@
+import os
+import cv2
+import glob
+import torch
+import shutil
+import numpy as np
+from tqdm import tqdm
+from util.reverse2original import reverse2wholeimage
+import moviepy.editor as mp
+from moviepy.editor import AudioFileClip, VideoFileClip
+from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
+import time
+from util.add_watermark import watermark_image
+from util.norm import SpecificNorm
+import torch.nn.functional as F
+from parsing_model.model import BiSeNet
+
+def _totensor(array):
+ tensor = torch.from_numpy(array)
+ img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
+ return img.float().div(255)
+
+def video_swap(video_path, id_vetor,specific_person_id_nonorm,id_thres, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224, no_simswaplogo = False,use_mask =False):
+ video_forcheck = VideoFileClip(video_path)
+ if video_forcheck.audio is None:
+ no_audio = True
+ else:
+ no_audio = False
+
+ del video_forcheck
+
+ if not no_audio:
+ video_audio_clip = AudioFileClip(video_path)
+
+ video = cv2.VideoCapture(video_path)
+ logoclass = watermark_image('./simswaplogo/simswaplogo.png')
+ ret = True
+ frame_index = 0
+
+ frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+
+ # video_WIDTH = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
+
+ # video_HEIGHT = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+ fps = video.get(cv2.CAP_PROP_FPS)
+ if os.path.exists(temp_results_dir):
+ shutil.rmtree(temp_results_dir)
+
+ spNorm =SpecificNorm()
+ mse = torch.nn.MSELoss().cuda()
+
+ if use_mask:
+ n_classes = 19
+ net = BiSeNet(n_classes=n_classes)
+ net.cuda()
+ save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth')
+ net.load_state_dict(torch.load(save_pth))
+ net.eval()
+ else:
+ net =None
+
+ # while ret:
+ for frame_index in tqdm(range(frame_count)):
+ ret, frame = video.read()
+ if ret:
+ detect_results = detect_model.get(frame,crop_size)
+
+ if detect_results is not None:
+ # print(frame_index)
+ if not os.path.exists(temp_results_dir):
+ os.mkdir(temp_results_dir)
+ frame_align_crop_list = detect_results[0]
+ frame_mat_list = detect_results[1]
+
+ id_compare_values = []
+ frame_align_crop_tenor_list = []
+ for frame_align_crop in frame_align_crop_list:
+
+ # BGR TO RGB
+ # frame_align_crop_RGB = frame_align_crop[...,::-1]
+
+ frame_align_crop_tenor = _totensor(cv2.cvtColor(frame_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
+
+ frame_align_crop_tenor_arcnorm = spNorm(frame_align_crop_tenor)
+ frame_align_crop_tenor_arcnorm_downsample = F.interpolate(frame_align_crop_tenor_arcnorm, size=(112,112))
+ frame_align_crop_crop_id_nonorm = swap_model.netArc(frame_align_crop_tenor_arcnorm_downsample)
+
+ id_compare_values.append(mse(frame_align_crop_crop_id_nonorm,specific_person_id_nonorm).detach().cpu().numpy())
+ frame_align_crop_tenor_list.append(frame_align_crop_tenor)
+ id_compare_values_array = np.array(id_compare_values)
+ min_index = np.argmin(id_compare_values_array)
+ min_value = id_compare_values_array[min_index]
+ if min_value < id_thres:
+ swap_result = swap_model(None, frame_align_crop_tenor_list[min_index], id_vetor, None, True)[0]
+
+ reverse2wholeimage([frame_align_crop_tenor_list[min_index]], [swap_result], [frame_mat_list[min_index]], crop_size, frame, logoclass,\
+ os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo,pasring_model =net,use_mask= use_mask, norm = spNorm)
+ else:
+ if not os.path.exists(temp_results_dir):
+ os.mkdir(temp_results_dir)
+ frame = frame.astype(np.uint8)
+ if not no_simswaplogo:
+ frame = logoclass.apply_frames(frame)
+ cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame)
+
+ else:
+ if not os.path.exists(temp_results_dir):
+ os.mkdir(temp_results_dir)
+ frame = frame.astype(np.uint8)
+ if not no_simswaplogo:
+ frame = logoclass.apply_frames(frame)
+ cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame)
+ else:
+ break
+
+ video.release()
+
+ # image_filename_list = []
+ path = os.path.join(temp_results_dir,'*.jpg')
+ image_filenames = sorted(glob.glob(path))
+
+ clips = ImageSequenceClip(image_filenames,fps = fps)
+
+ if not no_audio:
+ clips = clips.set_audio(video_audio_clip)
+
+
+ clips.write_videofile(save_path,audio_codec='aac')
+
diff --git a/SimSwap/util/visualizer.py b/SimSwap/util/visualizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..584ac45db8b613ac439592f1dc1af0ff1fc44632
--- /dev/null
+++ b/SimSwap/util/visualizer.py
@@ -0,0 +1,131 @@
+import numpy as np
+import os
+import ntpath
+import time
+from . import util
+from . import html
+import scipy.misc
+try:
+ from StringIO import StringIO # Python 2.7
+except ImportError:
+ from io import BytesIO # Python 3.x
+
+class Visualizer():
+ def __init__(self, opt):
+ # self.opt = opt
+ self.tf_log = opt.tf_log
+ self.use_html = opt.isTrain and not opt.no_html
+ self.win_size = opt.display_winsize
+ self.name = opt.name
+ if self.tf_log:
+ import tensorflow as tf
+ self.tf = tf
+ self.log_dir = os.path.join(opt.checkpoints_dir, opt.name, 'logs')
+ self.writer = tf.summary.FileWriter(self.log_dir)
+
+ if self.use_html:
+ self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')
+ self.img_dir = os.path.join(self.web_dir, 'images')
+ print('create web directory %s...' % self.web_dir)
+ util.mkdirs([self.web_dir, self.img_dir])
+ self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
+ with open(self.log_name, "a") as log_file:
+ now = time.strftime("%c")
+ log_file.write('================ Training Loss (%s) ================\n' % now)
+
+ # |visuals|: dictionary of images to display or save
+ def display_current_results(self, visuals, epoch, step):
+ if self.tf_log: # show images in tensorboard output
+ img_summaries = []
+ for label, image_numpy in visuals.items():
+ # Write the image to a string
+ try:
+ s = StringIO()
+ except:
+ s = BytesIO()
+ scipy.misc.toimage(image_numpy).save(s, format="jpeg")
+ # Create an Image object
+ img_sum = self.tf.Summary.Image(encoded_image_string=s.getvalue(), height=image_numpy.shape[0], width=image_numpy.shape[1])
+ # Create a Summary value
+ img_summaries.append(self.tf.Summary.Value(tag=label, image=img_sum))
+
+ # Create and write Summary
+ summary = self.tf.Summary(value=img_summaries)
+ self.writer.add_summary(summary, step)
+
+ if self.use_html: # save images to a html file
+ for label, image_numpy in visuals.items():
+ if isinstance(image_numpy, list):
+ for i in range(len(image_numpy)):
+ img_path = os.path.join(self.img_dir, 'epoch%.3d_%s_%d.jpg' % (epoch, label, i))
+ util.save_image(image_numpy[i], img_path)
+ else:
+ img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.jpg' % (epoch, label))
+ util.save_image(image_numpy, img_path)
+
+ # update website
+ webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, refresh=30)
+ for n in range(epoch, 0, -1):
+ webpage.add_header('epoch [%d]' % n)
+ ims = []
+ txts = []
+ links = []
+
+ for label, image_numpy in visuals.items():
+ if isinstance(image_numpy, list):
+ for i in range(len(image_numpy)):
+ img_path = 'epoch%.3d_%s_%d.jpg' % (n, label, i)
+ ims.append(img_path)
+ txts.append(label+str(i))
+ links.append(img_path)
+ else:
+ img_path = 'epoch%.3d_%s.jpg' % (n, label)
+ ims.append(img_path)
+ txts.append(label)
+ links.append(img_path)
+ if len(ims) < 10:
+ webpage.add_images(ims, txts, links, width=self.win_size)
+ else:
+ num = int(round(len(ims)/2.0))
+ webpage.add_images(ims[:num], txts[:num], links[:num], width=self.win_size)
+ webpage.add_images(ims[num:], txts[num:], links[num:], width=self.win_size)
+ webpage.save()
+
+ # errors: dictionary of error labels and values
+ def plot_current_errors(self, errors, step):
+ if self.tf_log:
+ for tag, value in errors.items():
+ summary = self.tf.Summary(value=[self.tf.Summary.Value(tag=tag, simple_value=value)])
+ self.writer.add_summary(summary, step)
+
+ # errors: same format as |errors| of plotCurrentErrors
+ def print_current_errors(self, epoch, i, errors, t):
+ message = '(epoch: %d, iters: %d, time: %.3f) ' % (epoch, i, t)
+ for k, v in errors.items():
+ if v != 0:
+ message += '%s: %.3f ' % (k, v)
+
+ print(message)
+ with open(self.log_name, "a") as log_file:
+ log_file.write('%s\n' % message)
+
+ # save image to the disk
+ def save_images(self, webpage, visuals, image_path):
+ image_dir = webpage.get_image_dir()
+ short_path = ntpath.basename(image_path[0])
+ name = os.path.splitext(short_path)[0]
+
+ webpage.add_header(name)
+ ims = []
+ txts = []
+ links = []
+
+ for label, image_numpy in visuals.items():
+ image_name = '%s_%s.jpg' % (name, label)
+ save_path = os.path.join(image_dir, image_name)
+ util.save_image(image_numpy, save_path)
+
+ ims.append(image_name)
+ txts.append(label)
+ links.append(image_name)
+ webpage.add_images(ims, txts, links, width=self.win_size)
+
+
+
+## Attention
+***This project is for technical and academic use only. Please do not apply it to illegal and unethical scenarios.***
+
+***In the event of violation of the legal and ethical requirements of the user's country or region, this code repository is exempt from liability***
+
+***Please do not ignore the content at the end of this README!***
+
+If you find this project useful, please star it. It is the greatest appreciation of our work.
+
+## Top News 
+
+**`2023-09-26`**: We fixed bugs in colab!
+
+**`2023-04-25`**: We fixed the "AttributeError: 'SGD' object has no attribute 'defaults' now" bug. If you have already downloaded **arcface_checkpoint.tar**, please **download it again**. Also, you also need to update the scripts in ```./models/```.
+
+**`2022-04-21`**: For resource limited users, we provide the cropped VGGFace2-224 dataset [[Google Driver] VGGFace2-224 (10.8G)](https://drive.google.com/file/d/19pWvdEHS-CEG6tW3PdxdtZ5QEymVjImc/view?usp=sharing) [[Baidu Driver]](https://pan.baidu.com/s/1OiwLJHVBSYB4AY2vEcfN0A) [Password: lrod].
+
+**`2022-04-20`**: Training scripts are now available. We highly recommend that you guys train the simswap model with our released high quality dataset [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ).
+
+**`2021-11-24`**: We have trained a beta version of ***SimSwap-HQ*** on [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ) and open sourced the checkpoint of this model (if you think the Simswap 512 is cool, please star our [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ) repo). Please don’t forget to go to [Preparation](./docs/guidance/preparation.md) and [Inference for image or video face swapping](./docs/guidance/usage.md) to check the latest set up.
+
+**`2021-11-23`**: The google drive link of [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ) is released.
+
+**`2021-11-17`**: We released a high resolution face dataset [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ) and the method to generate this dataset. This dataset is for research purpose.
+
+**`2021-08-30`**: Docker has been supported, please refer [here](https://replicate.ai/neuralchen/simswap-image) for details.
+
+**`2021-08-17`**: We have updated the [Preparation](./docs/guidance/preparation.md), The main change is that the gpu version of onnx is now installed by default, Now the time to process a video is greatly reduced.
+
+**`2021-07-19`**: ***Obvious border abruptness has been resolved***. We add the ability to using mask and upgrade the old algorithm for better visual effect, please go to [Inference for image or video face swapping](./docs/guidance/usage.md) for details. Please don’t forget to go to [Preparation](./docs/guidance/preparation.md) to check the latest set up. (Thanks for the help from [@woctezuma](https://github.com/woctezuma) and [@instant-high](https://github.com/instant-high))
+
+## The first open source high resolution dataset for face swapping!!!
+## High Resolution Dataset [VGGFace2-HQ](https://github.com/NNNNAI/VGGFace2-HQ)
+
+[](https://github.com/NNNNAI/VGGFace2-HQ)
+
+
+
+
+## Dependencies
+- python3.6+
+- pytorch1.5+
+- torchvision
+- opencv
+- pillow
+- numpy
+- imageio
+- moviepy
+- insightface
+- ***timm==0.5.4***
+
+## Training
+
+[Preparation](./docs/guidance/preparation.md)
+
+The training script is slightly different from the original version, e.g., we replace the patch discriminator with the projected discriminator, which saves a lot of hardware overhead and achieves slightly better results.
+
+In order to ensure the normal training, the batch size must be greater than 1.
+
+Friendly reminder, due to the difference in training settings, the user-trained model will have subtle differences in visual effects from the pre-trained model we provide.
+
+- Train 224 models with VGGFace2 224*224 [[Google Driver] VGGFace2-224 (10.8G)](https://drive.google.com/file/d/19pWvdEHS-CEG6tW3PdxdtZ5QEymVjImc/view?usp=sharing) [[Baidu Driver] ](https://pan.baidu.com/s/1OiwLJHVBSYB4AY2vEcfN0A) [Password: lrod]
+
+For faster convergence and better results, a large batch size (more than 16) is recommended!
+
+***We recommend training more than 400K iterations (batch size is 16), 600K~800K will be better, more iterations will not be recommended.***
+
+
+```
+python train.py --name simswap224_test --batchSize 8 --gpu_ids 0 --dataset /path/to/VGGFace2HQ --Gdeep False
+```
+
+[Colab demo for training 224 model][ 
+
+
+
+
+
+
+
+
+### About watermark of simswap logo
+The above example command lines are to add the simswap logo as the watermark by default. After our discussion, we have added a hyper parameter to control whether to remove watermark.
+
+The usage of removing the watermark is to add an argument: "***--no_simswaplogo***" to the command line, take the command line of "Swap all faces within one image" as an example, the following command line can get the result without watermark:
+```
+python test_wholeimage_swapmulti.py --no_simswaplogo --crop_size 224 --use_mask --name people --Arc_path arcface_model/arcface_checkpoint.tar --pic_a_path ./demo_file/Iron_man.jpg --pic_b_path ./demo_file/multi_people.jpg --output_path ./output/
+```
+### About using mask for better result
+We provide two methods to paste the face back to the original image after changing the face: Using mask or using bounding box. At present, the effect of using mask is the best. All the above code examples are using mask. If you want to use the bounding box, you only need to remove the --use_mask in the code example.
+Difference between using mask and not using mask can be found [here](https://imgsli.com/NjE3OTA).
+
+### Difference between single face swapping and all face swapping are shown below.
+
+
+
+
+
+### Parameters
+| Parameters | Function |
+| :---- | :---- |
+| --name | The SimSwap training logs name |
+| --pic_a_path | Path of image with the target face |
+| --pic_b_path | Path of image with the source face to swap |
+| --pic_specific_path | Path of image with the specific face to be swapped |
+|--multisepcific_dir |Path of image folder for multi specific face swapping|
+| --video_path | Path of video with the source face to swap |
+| --temp_path | Path to store intermediate files |
+| --output_path | Path of directory to store the face swapping result |
+| --no_simswaplogo |The hyper parameter to control whether to remove watermark |
+| --use_mask |The hyper parameter to control whether to use face parsing for the better visual effects(I recommend to use)|
+
+### Note
+We expect users to have GPU with at least 3G memory.the For those who do not, we will provide Colab Notebook implementation in the future.
diff --git a/SimSwap/docs/img/LRGT_201110059_201110091.webp b/SimSwap/docs/img/LRGT_201110059_201110091.webp
new file mode 100644
index 0000000000000000000000000000000000000000..72b66b6e3c00109883009bcce36a4a222b7e9c8a
--- /dev/null
+++ b/SimSwap/docs/img/LRGT_201110059_201110091.webp
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5ee03f9b50cbd670775d0d35d069d80faf6caa0852efaa4c56056fa2dcfb32ec
+size 1103968
diff --git a/SimSwap/docs/img/anni.webp b/SimSwap/docs/img/anni.webp
new file mode 100644
index 0000000000000000000000000000000000000000..6d92f047fb15392cbe60c79ef1a5546f7aa525d2
--- /dev/null
+++ b/SimSwap/docs/img/anni.webp
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b24ed6d04bda9064989d3a5cb175fa1b5a47f19c60560f8b8fc48f542122d9a
+size 4607822
diff --git a/SimSwap/docs/img/chenglong.webp b/SimSwap/docs/img/chenglong.webp
new file mode 100644
index 0000000000000000000000000000000000000000..94b73f08bdae2b7bf99f938ecd189b68e2fd3513
--- /dev/null
+++ b/SimSwap/docs/img/chenglong.webp
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1225e4e305aea3f397b026cdb21cd65f9153f03cfaceb48878a1f1ed383801fc
+size 4609164
diff --git a/SimSwap/docs/img/girl2-RGB.png b/SimSwap/docs/img/girl2-RGB.png
new file mode 100644
index 0000000000000000000000000000000000000000..a6caec4a018dcef972203606e0bb17f2f7651bd3
--- /dev/null
+++ b/SimSwap/docs/img/girl2-RGB.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ee8c1f81db555f4744f9f1d67b8ce2a90dfe4506bce2c6c295ba01531f40f3b
+size 1571819
diff --git a/SimSwap/docs/img/girl2.gif b/SimSwap/docs/img/girl2.gif
new file mode 100644
index 0000000000000000000000000000000000000000..f533a6eeb88fecb6b2e3000e4fa781b2ce8fe9b5
--- /dev/null
+++ b/SimSwap/docs/img/girl2.gif
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a357007c63e82866caaf95b67d20f5fc5cf4a74c55b9f6f160db5739372145b1
+size 16381792
diff --git a/SimSwap/docs/img/id/Iron_man.jpg b/SimSwap/docs/img/id/Iron_man.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..77a0f20b5d12dec5bd14488c80a36faf0a5d0778
Binary files /dev/null and b/SimSwap/docs/img/id/Iron_man.jpg differ
diff --git a/SimSwap/docs/img/id/anni.jpg b/SimSwap/docs/img/id/anni.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..ff42e9ecd0754b874b9da2eeb106da3f14fceac1
Binary files /dev/null and b/SimSwap/docs/img/id/anni.jpg differ
diff --git a/SimSwap/docs/img/id/chenglong.jpg b/SimSwap/docs/img/id/chenglong.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..8bbe4eafa6f29372bfa92694e0ce87943295494b
Binary files /dev/null and b/SimSwap/docs/img/id/chenglong.jpg differ
diff --git a/SimSwap/docs/img/id/wuyifan.png b/SimSwap/docs/img/id/wuyifan.png
new file mode 100644
index 0000000000000000000000000000000000000000..919d4ea916bc99574c958dc3e9afc3a015f5e499
--- /dev/null
+++ b/SimSwap/docs/img/id/wuyifan.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c0162a0446c27d0477796b84af948918b4af4056862df78f76391e1b8475bb7
+size 244429
diff --git a/SimSwap/docs/img/id/zhoujielun.jpg b/SimSwap/docs/img/id/zhoujielun.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..795d14b3dc8c78e597844ee5e9830f8aac233a1f
Binary files /dev/null and b/SimSwap/docs/img/id/zhoujielun.jpg differ
diff --git a/SimSwap/docs/img/id/zhuyin.jpg b/SimSwap/docs/img/id/zhuyin.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..f36434ba486dbf6434de176fad9e63335e039973
--- /dev/null
+++ b/SimSwap/docs/img/id/zhuyin.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4067296af9f24912b81717599d8e719eac9801773ca81c8fd3c3de61fcf0ac7e
+size 234685
diff --git a/SimSwap/docs/img/logo.png b/SimSwap/docs/img/logo.png
new file mode 100644
index 0000000000000000000000000000000000000000..7b889044fa62a5d6b18c04438738b24f77c7b6b7
Binary files /dev/null and b/SimSwap/docs/img/logo.png differ
diff --git a/SimSwap/docs/img/logo1.png b/SimSwap/docs/img/logo1.png
new file mode 100644
index 0000000000000000000000000000000000000000..36b993e38efa9382a9c20a2a03c69bc3002edb48
Binary files /dev/null and b/SimSwap/docs/img/logo1.png differ
diff --git a/SimSwap/docs/img/logo2.png b/SimSwap/docs/img/logo2.png
new file mode 100644
index 0000000000000000000000000000000000000000..884da60da8bc840878a9dde44f8bf4f838d7b8d2
Binary files /dev/null and b/SimSwap/docs/img/logo2.png differ
diff --git a/SimSwap/docs/img/mama_mask_short.webp b/SimSwap/docs/img/mama_mask_short.webp
new file mode 100644
index 0000000000000000000000000000000000000000..6fc8e6cce18de7788de188f3f37cd17428c3e9f4
--- /dev/null
+++ b/SimSwap/docs/img/mama_mask_short.webp
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0fcae8825b4432c707213ba1c4bb12a3558f21dc4a5daa33093236a64c4068aa
+size 8420998
diff --git a/SimSwap/docs/img/mama_mask_wuyifan_short.webp b/SimSwap/docs/img/mama_mask_wuyifan_short.webp
new file mode 100644
index 0000000000000000000000000000000000000000..c74fb237810a1dfa28936d2d0c1b9af9b02591ab
--- /dev/null
+++ b/SimSwap/docs/img/mama_mask_wuyifan_short.webp
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ce1a1268f53247815a40ced34d71a48f297a813f1c46a72fd60a3ac6db1873f
+size 8447672
diff --git a/SimSwap/docs/img/multi_face_comparison.png b/SimSwap/docs/img/multi_face_comparison.png
new file mode 100644
index 0000000000000000000000000000000000000000..6e7c6908bf63089e437154c7c75deab64fc42ce8
--- /dev/null
+++ b/SimSwap/docs/img/multi_face_comparison.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5094cc0e43097f1f36ce56cf97f8642fb230d380402cb4bdd18698f4d427a72e
+size 9173583
diff --git a/SimSwap/docs/img/new.gif b/SimSwap/docs/img/new.gif
new file mode 100644
index 0000000000000000000000000000000000000000..a5fe07d868f5ce8f163c218605c616159540fb4d
Binary files /dev/null and b/SimSwap/docs/img/new.gif differ
diff --git a/SimSwap/docs/img/nrsig.png b/SimSwap/docs/img/nrsig.png
new file mode 100644
index 0000000000000000000000000000000000000000..f5c562c16b837046ae2a201ce94bcdea65e91b8a
--- /dev/null
+++ b/SimSwap/docs/img/nrsig.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:410c7d637466e05c9994478956ab92604514b84c3c5005984a3e8ab8ced9bb31
+size 1597563
diff --git a/SimSwap/docs/img/result_whole_swap_multispecific_512.jpg b/SimSwap/docs/img/result_whole_swap_multispecific_512.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..a0b8aeb1345fe9c7783e9da1fb4a4492abbb018e
--- /dev/null
+++ b/SimSwap/docs/img/result_whole_swap_multispecific_512.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8cf2075ddb5d029d4c5d41f9c4adabada3fbf94268e06e1692e8000f2b42e1bf
+size 342746
diff --git a/SimSwap/docs/img/results1.PNG b/SimSwap/docs/img/results1.PNG
new file mode 100644
index 0000000000000000000000000000000000000000..4466b522c850b46a30b7e1eadee57463fa0aecef
--- /dev/null
+++ b/SimSwap/docs/img/results1.PNG
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5bcdf0796760d0a0e1b1dbcb9ab1960739e20461de6aaf73e16c8a2b17a1c53
+size 1626582
diff --git a/SimSwap/docs/img/title.png b/SimSwap/docs/img/title.png
new file mode 100644
index 0000000000000000000000000000000000000000..004ac600ce50761102fbc22787fc152c296afb89
--- /dev/null
+++ b/SimSwap/docs/img/title.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94e2dc3707e9ec846fe89bfb90be7586057a71de6c33872e9703592057a796f7
+size 7854435
diff --git a/SimSwap/docs/img/total.PNG b/SimSwap/docs/img/total.PNG
new file mode 100644
index 0000000000000000000000000000000000000000..0133fb9e66ac213669b8aadc5add96c3a92de886
--- /dev/null
+++ b/SimSwap/docs/img/total.PNG
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c28d4af9d932f34aa073a4eab6b11b60b824ad231da114c72b7405907d083bbf
+size 3169203
diff --git a/SimSwap/docs/img/vggface2_hq_compare.png b/SimSwap/docs/img/vggface2_hq_compare.png
new file mode 100644
index 0000000000000000000000000000000000000000..b91b46b60b067f8e7c84e52d39006ffde4278cc0
--- /dev/null
+++ b/SimSwap/docs/img/vggface2_hq_compare.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d9dd537d27099fe6a260552b3f512d2adc88e24f931e5f91029bffc08077b1c2
+size 1746262
diff --git a/SimSwap/docs/img/video.webp b/SimSwap/docs/img/video.webp
new file mode 100644
index 0000000000000000000000000000000000000000..e5eedf38873fc57e87509a54055f180539233ba3
--- /dev/null
+++ b/SimSwap/docs/img/video.webp
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c74ef3c0e7b8cbeab4ac63a6294aa012ad3c9659fb72745bd2ff6b2b2df82fc7
+size 3785104
diff --git a/SimSwap/docs/img/zhoujielun.webp b/SimSwap/docs/img/zhoujielun.webp
new file mode 100644
index 0000000000000000000000000000000000000000..5f9bb1c597d51bdfa956fd0addb91e4c8d054a72
--- /dev/null
+++ b/SimSwap/docs/img/zhoujielun.webp
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed89e5640aff5372bd301825d249d6460a2d2de88980071a49387a870ca4c018
+size 4609243
diff --git a/SimSwap/docs/img/zhuyin.webp b/SimSwap/docs/img/zhuyin.webp
new file mode 100644
index 0000000000000000000000000000000000000000..dc534ee7121cdf3a6f43c10ea815ba425b16004d
--- /dev/null
+++ b/SimSwap/docs/img/zhuyin.webp
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:434472dd31f93117976a1bd684ea9f7cfbf45c59fc52d35961e0e88a4ed11524
+size 4616881
diff --git a/SimSwap/docs/index.html b/SimSwap/docs/index.html
new file mode 100644
index 0000000000000000000000000000000000000000..2b48b1acc11e938e49edd753a312e34c8638c9fa
--- /dev/null
+++ b/SimSwap/docs/index.html
@@ -0,0 +1,345 @@
+
+
+
+
+
+
+
+
+
+
+
+
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+ 
+