From c77213d9198e0bbcb210b5e20031a30584d95377 Mon Sep 17 00:00:00 2001 From: "ryan.yy" Date: Tue, 10 Jan 2023 05:45:55 +0800 Subject: [PATCH] =?UTF-8?q?=E5=9B=BE=E5=83=8F=E6=8D=A2=E8=84=B8=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B=E4=B8=8AMaaS=20=20=20=20=20=20=20=20=20Link:=20https:?= =?UTF-8?q?//code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11347556?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- data/test/images/facefusion_template.jpg | 3 + data/test/images/facefusion_user.jpg | 3 + modelscope/metainfo.py | 2 + .../peppa_pig_face/LK/__init__.py | 0 .../cv/face_detection/peppa_pig_face/LK/lk.py | 97 +++ .../face_detection/peppa_pig_face/__init__.py | 0 .../peppa_pig_face/face_detector.py | 115 +++ .../peppa_pig_face/face_landmark.py | 154 ++++ .../cv/face_detection/peppa_pig_face/facer.py | 138 +++ .../models/cv/image_face_fusion/__init__.py | 20 + .../cv/image_face_fusion/facegan/__init__.py | 0 .../cv/image_face_fusion/facegan/gan_wrap.py | 93 +++ .../cv/image_face_fusion/facegan/model.py | 788 ++++++++++++++++++ .../image_face_fusion/facegan/op/__init__.py | 4 + .../facegan/op/conv2d_gradfix.py | 228 +++++ .../image_face_fusion/facegan/op/fused_act.py | 113 +++ .../image_face_fusion/facegan/op/upfirdn2d.py | 198 +++++ .../cv/image_face_fusion/facelib/__init__.py | 0 .../image_face_fusion/facelib/align_trans.py | 301 +++++++ .../facelib/matlab_cp2tform.py | 230 +++++ .../cv/image_face_fusion/image_face_fusion.py | 253 ++++++ .../cv/image_face_fusion/network/__init__.py | 0 .../cv/image_face_fusion/network/aad_layer.py | 99 +++ .../image_face_fusion/network/aei_flow_net.py | 251 ++++++ .../cv/image_face_fusion/network/bfm.py | 249 ++++++ .../image_face_fusion/network/dense_motion.py | 376 +++++++++ .../network/facerecon_model.py | 559 +++++++++++++ .../image_face_fusion/network/model_irse.py | 249 ++++++ .../cv/image_face_fusion/network/ops.py | 211 +++++ modelscope/pipelines/builder.py | 4 + .../cv/image_face_fusion_pipeline.py | 68 ++ modelscope/utils/constant.py | 2 +- tests/pipelines/test_image_face_fusion.py | 59 ++ 33 files changed, 4866 insertions(+), 1 deletion(-) create mode 100644 data/test/images/facefusion_template.jpg create mode 100644 data/test/images/facefusion_user.jpg create mode 100644 modelscope/models/cv/face_detection/peppa_pig_face/LK/__init__.py create mode 100644 modelscope/models/cv/face_detection/peppa_pig_face/LK/lk.py create mode 100644 modelscope/models/cv/face_detection/peppa_pig_face/__init__.py create mode 100644 modelscope/models/cv/face_detection/peppa_pig_face/face_detector.py create mode 100644 modelscope/models/cv/face_detection/peppa_pig_face/face_landmark.py create mode 100644 modelscope/models/cv/face_detection/peppa_pig_face/facer.py create mode 100644 modelscope/models/cv/image_face_fusion/__init__.py create mode 100644 modelscope/models/cv/image_face_fusion/facegan/__init__.py create mode 100644 modelscope/models/cv/image_face_fusion/facegan/gan_wrap.py create mode 100644 modelscope/models/cv/image_face_fusion/facegan/model.py create mode 100644 modelscope/models/cv/image_face_fusion/facegan/op/__init__.py create mode 100644 modelscope/models/cv/image_face_fusion/facegan/op/conv2d_gradfix.py create mode 100644 modelscope/models/cv/image_face_fusion/facegan/op/fused_act.py create mode 100644 modelscope/models/cv/image_face_fusion/facegan/op/upfirdn2d.py create mode 100644 modelscope/models/cv/image_face_fusion/facelib/__init__.py create mode 100644 modelscope/models/cv/image_face_fusion/facelib/align_trans.py create mode 100644 modelscope/models/cv/image_face_fusion/facelib/matlab_cp2tform.py create mode 100644 modelscope/models/cv/image_face_fusion/image_face_fusion.py create mode 100644 modelscope/models/cv/image_face_fusion/network/__init__.py create mode 100644 modelscope/models/cv/image_face_fusion/network/aad_layer.py create mode 100644 modelscope/models/cv/image_face_fusion/network/aei_flow_net.py create mode 100644 modelscope/models/cv/image_face_fusion/network/bfm.py create mode 100644 modelscope/models/cv/image_face_fusion/network/dense_motion.py create mode 100644 modelscope/models/cv/image_face_fusion/network/facerecon_model.py create mode 100644 modelscope/models/cv/image_face_fusion/network/model_irse.py create mode 100644 modelscope/models/cv/image_face_fusion/network/ops.py create mode 100644 modelscope/pipelines/cv/image_face_fusion_pipeline.py create mode 100644 tests/pipelines/test_image_face_fusion.py diff --git a/data/test/images/facefusion_template.jpg b/data/test/images/facefusion_template.jpg new file mode 100644 index 00000000..60a8a6cb --- /dev/null +++ b/data/test/images/facefusion_template.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5e11e3558040246fc6d84bf87afdb016228172893f475f843dedbdcda5092a3d +size 181713 diff --git a/data/test/images/facefusion_user.jpg b/data/test/images/facefusion_user.jpg new file mode 100644 index 00000000..b8d0b5df --- /dev/null +++ b/data/test/images/facefusion_user.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e82e688d2eb2755ceb0b0051d7129f6e94e6e5fe57f68727e41cd0c1e909b89c +size 11143 diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index bda6f41b..4ad383bd 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -69,6 +69,7 @@ class Models(object): rcp_sceneflow_estimation = 'rcp-sceneflow-estimation' image_casmvs_depth_estimation = 'image-casmvs-depth-estimation' ddcolor = 'ddcolor' + image_face_fusion = 'image-face-fusion' # EasyCV models yolox = 'YOLOX' @@ -275,6 +276,7 @@ class Pipelines(object): pointcloud_sceneflow_estimation = 'pointcloud-sceneflow-estimation' image_multi_view_depth_estimation = 'image-multi-view-depth-estimation' ddcolor_image_colorization = 'ddcolor-image-colorization' + image_face_fusion = 'image-face-fusion' # nlp tasks automatic_post_editing = 'automatic-post-editing' diff --git a/modelscope/models/cv/face_detection/peppa_pig_face/LK/__init__.py b/modelscope/models/cv/face_detection/peppa_pig_face/LK/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/face_detection/peppa_pig_face/LK/lk.py b/modelscope/models/cv/face_detection/peppa_pig_face/LK/lk.py new file mode 100644 index 00000000..3b4a670c --- /dev/null +++ b/modelscope/models/cv/face_detection/peppa_pig_face/LK/lk.py @@ -0,0 +1,97 @@ +# The implementation here is modified based on InsightFace_Pytorch, originally Apache License and publicly available +# at https://github.com/610265158/Peppa_Pig_Face_Engine +import numpy as np + + +class GroupTrack(): + + def __init__(self): + self.old_frame = None + self.previous_landmarks_set = None + self.with_landmark = True + self.thres = 1 + self.alpha = 0.95 + self.iou_thres = 0.5 + + def calculate(self, img, current_landmarks_set): + if self.previous_landmarks_set is None: + self.previous_landmarks_set = current_landmarks_set + result = current_landmarks_set + else: + previous_lm_num = self.previous_landmarks_set.shape[0] + if previous_lm_num == 0: + self.previous_landmarks_set = current_landmarks_set + result = current_landmarks_set + return result + else: + result = [] + for i in range(current_landmarks_set.shape[0]): + not_in_flag = True + for j in range(previous_lm_num): + if self.iou(current_landmarks_set[i], + self.previous_landmarks_set[j] + ) > self.iou_thres: + result.append( + self.smooth(current_landmarks_set[i], + self.previous_landmarks_set[j])) + not_in_flag = False + break + if not_in_flag: + result.append(current_landmarks_set[i]) + + result = np.array(result) + self.previous_landmarks_set = result + + return result + + def iou(self, p_set0, p_set1): + rec1 = [ + np.min(p_set0[:, 0]), + np.min(p_set0[:, 1]), + np.max(p_set0[:, 0]), + np.max(p_set0[:, 1]) + ] + rec2 = [ + np.min(p_set1[:, 0]), + np.min(p_set1[:, 1]), + np.max(p_set1[:, 0]), + np.max(p_set1[:, 1]) + ] + + # computing area of each rectangles + S_rec1 = (rec1[2] - rec1[0]) * (rec1[3] - rec1[1]) + S_rec2 = (rec2[2] - rec2[0]) * (rec2[3] - rec2[1]) + + # computing the sum_area + sum_area = S_rec1 + S_rec2 + + # find the each edge of intersect rectangle + x1 = max(rec1[0], rec2[0]) + y1 = max(rec1[1], rec2[1]) + x2 = min(rec1[2], rec2[2]) + y2 = min(rec1[3], rec2[3]) + + # judge if there is an intersect + intersect = max(0, x2 - x1) * max(0, y2 - y1) + + iou = intersect / (sum_area - intersect) + return iou + + def smooth(self, now_landmarks, previous_landmarks): + result = [] + for i in range(now_landmarks.shape[0]): + x = now_landmarks[i][0] - previous_landmarks[i][0] + y = now_landmarks[i][1] - previous_landmarks[i][1] + dis = np.sqrt(np.square(x) + np.square(y)) + if dis < self.thres: + result.append(previous_landmarks[i]) + else: + result.append( + self.do_moving_average(now_landmarks[i], + previous_landmarks[i])) + + return np.array(result) + + def do_moving_average(self, p_now, p_previous): + p = self.alpha * p_now + (1 - self.alpha) * p_previous + return p diff --git a/modelscope/models/cv/face_detection/peppa_pig_face/__init__.py b/modelscope/models/cv/face_detection/peppa_pig_face/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/face_detection/peppa_pig_face/face_detector.py b/modelscope/models/cv/face_detection/peppa_pig_face/face_detector.py new file mode 100644 index 00000000..4deff9bf --- /dev/null +++ b/modelscope/models/cv/face_detection/peppa_pig_face/face_detector.py @@ -0,0 +1,115 @@ +# The implementation here is modified based on InsightFace_Pytorch, originally Apache License and publicly available +# at https://github.com/610265158/Peppa_Pig_Face_Engine +import cv2 +import numpy as np +import tensorflow as tf + +if tf.__version__ >= '2.0': + tf = tf.compat.v1 + + +class FaceDetector: + + def __init__(self, dir): + + self.model_path = dir + '/detector.pb' + self.thres = 0.8 + self.input_shape = (512, 512, 3) + self.pixel_means = np.array([123., 116., 103.]) + + self._graph = tf.Graph() + + with self._graph.as_default(): + self._graph, self._sess = self.init_model(self.model_path) + + self.input_image = tf.get_default_graph().get_tensor_by_name( + 'tower_0/images:0') + self.training = tf.get_default_graph().get_tensor_by_name( + 'training_flag:0') + self.output_ops = [ + tf.get_default_graph().get_tensor_by_name('tower_0/boxes:0'), + tf.get_default_graph().get_tensor_by_name('tower_0/scores:0'), + tf.get_default_graph().get_tensor_by_name( + 'tower_0/num_detections:0'), + ] + + def __call__(self, image): + + image, scale_x, scale_y = self.preprocess( + image, + target_width=self.input_shape[1], + target_height=self.input_shape[0]) + + image = np.expand_dims(image, 0) + + boxes, scores, num_boxes = self._sess.run( + self.output_ops, + feed_dict={ + self.input_image: image, + self.training: False + }) + + num_boxes = num_boxes[0] + boxes = boxes[0][:num_boxes] + + scores = scores[0][:num_boxes] + + to_keep = scores > self.thres + boxes = boxes[to_keep] + scores = scores[to_keep] + + y1 = self.input_shape[0] / scale_y + x1 = self.input_shape[1] / scale_x + y2 = self.input_shape[0] / scale_y + x2 = self.input_shape[1] / scale_x + scaler = np.array([y1, x1, y2, x2], dtype='float32') + boxes = boxes * scaler + + scores = np.expand_dims(scores, 0).reshape([-1, 1]) + + for i in range(boxes.shape[0]): + boxes[i] = np.array( + [boxes[i][1], boxes[i][0], boxes[i][3], boxes[i][2]]) + return np.concatenate([boxes, scores], axis=1) + + def preprocess(self, image, target_height, target_width, label=None): + + h, w, c = image.shape + + bimage = np.zeros( + shape=[target_height, target_width, c], + dtype=image.dtype) + np.array( + self.pixel_means, dtype=image.dtype) + long_side = max(h, w) + + scale_x = scale_y = target_height / long_side + + image = cv2.resize(image, None, fx=scale_x, fy=scale_y) + + h_, w_, _ = image.shape + bimage[:h_, :w_, :] = image + + return bimage, scale_x, scale_y + + def init_model(self, *args): + pb_path = args[0] + + def init_pb(model_path): + config = tf.ConfigProto() + config.gpu_options.per_process_gpu_memory_fraction = 0.2 + compute_graph = tf.Graph() + compute_graph.as_default() + sess = tf.Session(config=config) + with tf.gfile.GFile(model_path, 'rb') as fid: + graph_def = tf.GraphDef() + graph_def.ParseFromString(fid.read()) + tf.import_graph_def(graph_def, name='') + + return (compute_graph, sess) + + model = init_pb(pb_path) + + graph = model[0] + sess = model[1] + + return graph, sess diff --git a/modelscope/models/cv/face_detection/peppa_pig_face/face_landmark.py b/modelscope/models/cv/face_detection/peppa_pig_face/face_landmark.py new file mode 100644 index 00000000..03a3b5b7 --- /dev/null +++ b/modelscope/models/cv/face_detection/peppa_pig_face/face_landmark.py @@ -0,0 +1,154 @@ +# The implementation here is modified based on InsightFace_Pytorch, originally Apache License and publicly available +# at https://github.com/610265158/Peppa_Pig_Face_Engine +import cv2 +import numpy as np +import tensorflow as tf + +if tf.__version__ >= '2.0': + tf = tf.compat.v1 + + +class FaceLandmark: + + def __init__(self, dir): + self.model_path = dir + '/keypoints.pb' + self.min_face = 60 + self.keypoint_num = 136 + self.pixel_means = np.array([123., 116., 103.]) + self.kp_extend_range = [0.2, 0.3] + self.kp_shape = (160, 160, 3) + + self._graph = tf.Graph() + + with self._graph.as_default(): + + self._graph, self._sess = self.init_model(self.model_path) + self.img_input = tf.get_default_graph().get_tensor_by_name( + 'tower_0/images:0') + self.embeddings = tf.get_default_graph().get_tensor_by_name( + 'tower_0/prediction:0') + self.training = tf.get_default_graph().get_tensor_by_name( + 'training_flag:0') + + self.landmark = self.embeddings[:, :self.keypoint_num] + self.headpose = self.embeddings[:, -7:-4] * 90. + self.state = tf.nn.sigmoid(self.embeddings[:, -4:]) + + def __call__(self, img, bboxes): + landmark_result = [] + state_result = [] + for i, bbox in enumerate(bboxes): + landmark, state = self._one_shot_run(img, bbox, i) + if landmark is not None: + landmark_result.append(landmark) + state_result.append(state) + return np.array(landmark_result), np.array(state_result) + + def simple_run(self, cropped_img): + with self._graph.as_default(): + + cropped_img = np.expand_dims(cropped_img, axis=0) + landmark, p, states = self._sess.run( + [self.landmark, self.headpose, self.state], + feed_dict={ + self.img_input: cropped_img, + self.training: False + }) + + return landmark, states + + def _one_shot_run(self, image, bbox, i): + + bbox_width = bbox[2] - bbox[0] + bbox_height = bbox[3] - bbox[1] + if (bbox_width <= self.min_face and bbox_height <= self.min_face): + return None, None + add = int(max(bbox_width, bbox_height)) + bimg = cv2.copyMakeBorder( + image, + add, + add, + add, + add, + borderType=cv2.BORDER_CONSTANT, + value=self.pixel_means) + bbox += add + + one_edge = (1 + 2 * self.kp_extend_range[0]) * bbox_width + center = [(bbox[0] + bbox[2]) // 2, (bbox[1] + bbox[3]) // 2] + + bbox[0] = center[0] - one_edge // 2 + bbox[1] = center[1] - one_edge // 2 + bbox[2] = center[0] + one_edge // 2 + bbox[3] = center[1] + one_edge // 2 + + bbox = bbox.astype(np.int) + crop_image = bimg[bbox[1]:bbox[3], bbox[0]:bbox[2], :] + h, w, _ = crop_image.shape + crop_image = cv2.resize(crop_image, + (self.kp_shape[1], self.kp_shape[0])) + crop_image = crop_image.astype(np.float32) + + keypoints, state = self.simple_run(crop_image) + + res = keypoints[0][:self.keypoint_num].reshape((-1, 2)) + res[:, 0] = res[:, 0] * w / self.kp_shape[1] + res[:, 1] = res[:, 1] * h / self.kp_shape[0] + + landmark = [] + for _index in range(res.shape[0]): + x_y = res[_index] + landmark.append([ + int(x_y[0] * self.kp_shape[0] + bbox[0] - add), + int(x_y[1] * self.kp_shape[1] + bbox[1] - add) + ]) + + landmark = np.array(landmark, np.float32) + + return landmark, state + + def init_model(self, *args): + + if len(args) == 1: + use_pb = True + pb_path = args[0] + else: + use_pb = False + meta_path = args[0] + restore_model_path = args[1] + + def ini_ckpt(): + graph = tf.Graph() + graph.as_default() + configProto = tf.ConfigProto() + configProto.gpu_options.allow_growth = True + sess = tf.Session(config=configProto) + # load_model(model_path, sess) + saver = tf.train.import_meta_graph(meta_path) + saver.restore(sess, restore_model_path) + + print('Model restored!') + return (graph, sess) + + def init_pb(model_path): + config = tf.ConfigProto() + config.gpu_options.per_process_gpu_memory_fraction = 0.2 + compute_graph = tf.Graph() + compute_graph.as_default() + sess = tf.Session(config=config) + with tf.gfile.GFile(model_path, 'rb') as fid: + graph_def = tf.GraphDef() + graph_def.ParseFromString(fid.read()) + tf.import_graph_def(graph_def, name='') + + return (compute_graph, sess) + + if use_pb: + model = init_pb(pb_path) + else: + model = ini_ckpt() + + graph = model[0] + sess = model[1] + + return graph, sess diff --git a/modelscope/models/cv/face_detection/peppa_pig_face/facer.py b/modelscope/models/cv/face_detection/peppa_pig_face/facer.py new file mode 100644 index 00000000..16039d83 --- /dev/null +++ b/modelscope/models/cv/face_detection/peppa_pig_face/facer.py @@ -0,0 +1,138 @@ +# The implementation here is modified based on InsightFace_Pytorch, originally Apache License and publicly available +# at https://github.com/610265158/Peppa_Pig_Face_Engine +import cv2 +import numpy as np + +from .face_detector import FaceDetector +from .face_landmark import FaceLandmark +from .LK.lk import GroupTrack + + +class FaceAna(): + + def __init__(self, model_dir): + self.face_detector = FaceDetector(model_dir) + self.face_landmark = FaceLandmark(model_dir) + self.trace = GroupTrack() + + self.track_box = None + self.previous_image = None + self.previous_box = None + + self.diff_thres = 5 + self.top_k = 10 + self.iou_thres = 0.5 + self.alpha = 0.3 + + def run(self, image): + + boxes = self.face_detector(image) + + if boxes.shape[0] > self.top_k: + boxes = self.sort(boxes) + + boxes_return = np.array(boxes) + landmarks, states = self.face_landmark(image, boxes) + + if 1: + track = [] + for i in range(landmarks.shape[0]): + track.append([ + np.min(landmarks[i][:, 0]), + np.min(landmarks[i][:, 1]), + np.max(landmarks[i][:, 0]), + np.max(landmarks[i][:, 1]) + ]) + tmp_box = np.array(track) + + self.track_box = self.judge_boxs(boxes_return, tmp_box) + + self.track_box, landmarks = self.sort_res(self.track_box, landmarks) + return self.track_box, landmarks, states + + def sort_res(self, bboxes, points): + area = [] + for bbox in bboxes: + bbox_width = bbox[2] - bbox[0] + bbox_height = bbox[3] - bbox[1] + area.append(bbox_height * bbox_width) + + area = np.array(area) + picked = area.argsort()[::-1] + sorted_bboxes = [bboxes[x] for x in picked] + sorted_points = [points[x] for x in picked] + return np.array(sorted_bboxes), np.array(sorted_points) + + def diff_frames(self, previous_frame, image): + if previous_frame is None: + return True + else: + _diff = cv2.absdiff(previous_frame, image) + diff = np.sum( + _diff) / previous_frame.shape[0] / previous_frame.shape[1] / 3. + return diff > self.diff_thres + + def sort(self, bboxes): + if self.top_k > 100: + return bboxes + area = [] + for bbox in bboxes: + + bbox_width = bbox[2] - bbox[0] + bbox_height = bbox[3] - bbox[1] + area.append(bbox_height * bbox_width) + + area = np.array(area) + + picked = area.argsort()[-self.top_k:][::-1] + sorted_bboxes = [bboxes[x] for x in picked] + return np.array(sorted_bboxes) + + def judge_boxs(self, previuous_bboxs, now_bboxs): + + def iou(rec1, rec2): + + S_rec1 = (rec1[2] - rec1[0]) * (rec1[3] - rec1[1]) + S_rec2 = (rec2[2] - rec2[0]) * (rec2[3] - rec2[1]) + + sum_area = S_rec1 + S_rec2 + + x1 = max(rec1[0], rec2[0]) + y1 = max(rec1[1], rec2[1]) + x2 = min(rec1[2], rec2[2]) + y2 = min(rec1[3], rec2[3]) + + intersect = max(0, x2 - x1) * max(0, y2 - y1) + + return intersect / (sum_area - intersect) + + if previuous_bboxs is None: + return now_bboxs + + result = [] + + for i in range(now_bboxs.shape[0]): + contain = False + for j in range(previuous_bboxs.shape[0]): + if iou(now_bboxs[i], previuous_bboxs[j]) > self.iou_thres: + result.append( + self.smooth(now_bboxs[i], previuous_bboxs[j])) + contain = True + break + if not contain: + result.append(now_bboxs[i]) + + return np.array(result) + + def smooth(self, now_box, previous_box): + + return self.do_moving_average(now_box[:4], previous_box[:4]) + + def do_moving_average(self, p_now, p_previous): + p = self.alpha * p_now + (1 - self.alpha) * p_previous + return p + + def reset(self): + self.track_box = None + self.previous_image = None + self.previous_box = None diff --git a/modelscope/models/cv/image_face_fusion/__init__.py b/modelscope/models/cv/image_face_fusion/__init__.py new file mode 100644 index 00000000..5c77e231 --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import TYPE_CHECKING + +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + from .image_face_fusion import ImageFaceFusion + +else: + _import_structure = {'image_face_fusion': ['ImageFaceFusion']} + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/modelscope/models/cv/image_face_fusion/facegan/__init__.py b/modelscope/models/cv/image_face_fusion/facegan/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/image_face_fusion/facegan/gan_wrap.py b/modelscope/models/cv/image_face_fusion/facegan/gan_wrap.py new file mode 100644 index 00000000..c46b17eb --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/facegan/gan_wrap.py @@ -0,0 +1,93 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +from torchvision import transforms + +from .model import FullGenerator + + +class GANWrap(object): + + def __init__(self, + model_path, + size=256, + channel_multiplier=1, + device='cpu'): + self.device = device + self.mfile = model_path + self.transform = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), + inplace=True), + ]) + self.batchSize = 2 + self.n_mlp = 8 + self.resolution = size + self.load_model(channel_multiplier) + + def load_model(self, channel_multiplier=2): + self.model = FullGenerator(self.resolution, 512, self.n_mlp, + channel_multiplier).to(self.device) + pretrained_dict = torch.load( + self.mfile, map_location=torch.device('cpu')) + self.model.load_state_dict(pretrained_dict) + self.model.eval() + + def process_tensor(self, img_t, return_face=True): + b, c, h, w = img_t.shape + img_t = F.interpolate(img_t, (self.resolution, self.resolution)) + + with torch.no_grad(): + out, __ = self.model(img_t) + + out = F.interpolate(out, (w, h)) + return out + + def process(self, ims, return_face=True): + res = [] + faces = [] + for i in range(0, len(ims), self.batchSize): + sizes = [] + imt = None + for im in ims[i:i + self.batchSize]: + sizes.append(im.shape[0]) + im = cv2.resize(im, (self.resolution, self.resolution)) + im_pil = Image.fromarray(im) + imt = self.img2tensor(im_pil) if imt is None else torch.cat( + (imt, self.img2tensor(im_pil)), dim=0) + + imt = torch.flip(imt, [1]) + with torch.no_grad(): + img_outs, __ = self.model(imt) + + for sz, img_out in zip(sizes, img_outs): + img = self.tensor2img(img_out) + if return_face: + faces.append(img) + img = cv2.resize(img, (sz, sz), interpolation=cv2.INTER_AREA) + res.append(img) + + return res, faces + + def img2tensor(self, img): + img_t = self.transform(img).to(self.device) + img_t = torch.unsqueeze(img_t, 0) + return img_t + + def tensor2img(self, image_tensor, bytes=255.0, imtype=np.uint8): + if image_tensor.dim() == 3: + image_numpy = image_tensor.cpu().float().numpy() + else: + image_numpy = image_tensor[0].cpu().float().numpy() + image_numpy = np.transpose(image_numpy, (1, 2, 0)) + image_numpy = image_numpy[:, :, ::-1] + image_numpy = np.clip( + image_numpy * np.asarray([0.5, 0.5, 0.5]) + + np.asarray([0.5, 0.5, 0.5]), 0, 1) + image_numpy = image_numpy * bytes + return image_numpy.astype(imtype) diff --git a/modelscope/models/cv/image_face_fusion/facegan/model.py b/modelscope/models/cv/image_face_fusion/facegan/model.py new file mode 100644 index 00000000..eb142779 --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/facegan/model.py @@ -0,0 +1,788 @@ +# The implementation is adopted from stylegan2-pytorch, +# made public available under the MIT License at https://github.com/rosinality/stylegan2-pytorch/blob/master/model.py +import math +import random + +import torch +from torch import nn +from torch.nn import functional as F + +from .op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d + +isconcat = True +sss = 2 if isconcat else 1 +ratio = 2 + + +class PixelNorm(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, input): + return input * torch.rsqrt( + torch.mean(input**2, dim=1, keepdim=True) + 1e-8) + + +def make_kernel(k): + k = torch.tensor(k, dtype=torch.float32) + + if k.ndim == 1: + k = k[None, :] * k[:, None] + + k /= k.sum() + + return k + + +class Upsample(nn.Module): + + def __init__(self, kernel, factor=2): + super().__init__() + + self.factor = factor + kernel = make_kernel(kernel) * (factor**2) + self.register_buffer('kernel', kernel) + + p = kernel.shape[0] - factor + + pad0 = (p + 1) // 2 + factor - 1 + pad1 = p // 2 + + self.pad = (pad0, pad1) + + def forward(self, input): + out = upfirdn2d( + input, self.kernel, up=self.factor, down=1, pad=self.pad) + + return out + + +class Downsample(nn.Module): + + def __init__(self, kernel, factor=2): + super().__init__() + + self.factor = factor + kernel = make_kernel(kernel) + self.register_buffer('kernel', kernel) + + p = kernel.shape[0] - factor + + pad0 = (p + 1) // 2 + pad1 = p // 2 + + self.pad = (pad0, pad1) + + def forward(self, input): + out = upfirdn2d( + input, self.kernel, up=1, down=self.factor, pad=self.pad) + + return out + + +class Blur(nn.Module): + + def __init__(self, kernel, pad, upsample_factor=1): + super().__init__() + + kernel = make_kernel(kernel) + + if upsample_factor > 1: + kernel = kernel * (upsample_factor**2) + + self.register_buffer('kernel', kernel) + + self.pad = pad + + def forward(self, input): + out = upfirdn2d(input, self.kernel, pad=self.pad) + + return out + + +class EqualConv2d(nn.Module): + + def __init__(self, + in_channel, + out_channel, + kernel_size, + stride=1, + padding=0, + bias=True): + super().__init__() + + self.weight = nn.Parameter( + torch.randn(out_channel, in_channel, kernel_size, kernel_size)) + self.scale = 1 / math.sqrt(in_channel * kernel_size**2) + + self.stride = stride + self.padding = padding + + if bias: + self.bias = nn.Parameter(torch.zeros(out_channel)) + + else: + self.bias = None + + def forward(self, input): + out = F.conv2d( + input, + self.weight * self.scale, + bias=self.bias, + stride=self.stride, + padding=self.padding, + ) + + return out + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},' + f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})' + ) + + +class EqualLinear(nn.Module): + + def __init__(self, + in_dim, + out_dim, + bias=True, + bias_init=0, + lr_mul=1, + activation=None): + super().__init__() + + self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul)) + + if bias: + self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init)) + + else: + self.bias = None + + self.activation = activation + + self.scale = (1 / math.sqrt(in_dim)) * lr_mul + self.lr_mul = lr_mul + + def forward(self, input): + if self.activation: + out = F.linear(input, self.weight * self.scale) + out = fused_leaky_relu(out, self.bias * self.lr_mul) + + else: + out = F.linear( + input, self.weight * self.scale, bias=self.bias * self.lr_mul) + + return out + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})' + ) + + +class ScaledLeakyReLU(nn.Module): + + def __init__(self, negative_slope=0.2): + super().__init__() + + self.negative_slope = negative_slope + + def forward(self, input): + out = F.leaky_relu(input, negative_slope=self.negative_slope) + + return out * math.sqrt(2) + + +class ModulatedConv2d(nn.Module): + + def __init__( + self, + in_channel, + out_channel, + kernel_size, + style_dim, + demodulate=True, + upsample=False, + downsample=False, + blur_kernel=[1, 3, 3, 1], + ): + super().__init__() + + self.eps = 1e-8 + self.kernel_size = kernel_size + self.in_channel = in_channel + self.out_channel = out_channel + self.upsample = upsample + self.downsample = downsample + + if upsample: + factor = 2 + p = (len(blur_kernel) - factor) - (kernel_size - 1) + pad0 = (p + 1) // 2 + factor - 1 + pad1 = p // 2 + 1 + + self.blur = Blur( + blur_kernel, pad=(pad0, pad1), upsample_factor=factor) + + if downsample: + factor = 2 + p = (len(blur_kernel) - factor) + (kernel_size - 1) + pad0 = (p + 1) // 2 + pad1 = p // 2 + + self.blur = Blur(blur_kernel, pad=(pad0, pad1)) + + fan_in = in_channel * kernel_size**2 + self.scale = 1 / math.sqrt(fan_in) + self.padding = kernel_size // 2 + + self.weight = nn.Parameter( + torch.randn(1, out_channel, in_channel, kernel_size, kernel_size)) + + self.modulation = EqualLinear(style_dim, in_channel, bias_init=1) + + self.demodulate = demodulate + + def __repr__(self): + return ( + f'{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, ' + f'upsample={self.upsample}, downsample={self.downsample})') + + def forward(self, input, style): + batch, in_channel, height, width = input.shape + + style = self.modulation(style).view(batch, 1, in_channel, 1, 1) + weight = self.scale * self.weight * style + + if self.demodulate: + demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8) + weight = weight * demod.view(batch, self.out_channel, 1, 1, 1) + + weight = weight.view(batch * self.out_channel, in_channel, + self.kernel_size, self.kernel_size) + + if self.upsample: + input = input.view(1, batch * in_channel, height, width) + weight = weight.view(batch, self.out_channel, in_channel, + self.kernel_size, self.kernel_size) + weight = weight.transpose(1, 2).reshape(batch * in_channel, + self.out_channel, + self.kernel_size, + self.kernel_size) + out = F.conv_transpose2d( + input, weight, padding=0, stride=2, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + out = self.blur(out) + + elif self.downsample: + input = self.blur(input) + _, _, height, width = input.shape + input = input.view(1, batch * in_channel, height, width) + out = F.conv2d(input, weight, padding=0, stride=2, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + + else: + input = input.view(1, batch * in_channel, height, width) + out = F.conv2d(input, weight, padding=self.padding, groups=batch) + _, _, height, width = out.shape + out = out.view(batch, self.out_channel, height, width) + + return out + + +class NoiseInjection(nn.Module): + + def __init__(self): + super().__init__() + + self.weight = nn.Parameter(torch.zeros(1)) + + def forward(self, image, noise=None): + + if noise is not None: + if isconcat: + return torch.cat((image, self.weight * noise), dim=1) # concat + return image + self.weight * noise + + if noise is None: + batch, _, height, width = image.shape + noise = image.new_empty(batch, 1, height, width).normal_() + + return image + self.weight * noise + + +class ConstantInput(nn.Module): + + def __init__(self, channel, size=4): + super().__init__() + + self.input = nn.Parameter(torch.randn(1, channel, size, size)) + + def forward(self, input): + batch = input.shape[0] + out = self.input.repeat(batch, 1, 1, 1) + + return out + + +class StyledConv(nn.Module): + + def __init__( + self, + in_channel, + out_channel, + kernel_size, + style_dim, + upsample=False, + blur_kernel=[1, 3, 3, 1], + demodulate=True, + ): + super().__init__() + + self.conv = ModulatedConv2d( + in_channel, + out_channel, + kernel_size, + style_dim, + upsample=upsample, + blur_kernel=blur_kernel, + demodulate=demodulate, + ) + + self.noise = NoiseInjection() + self.activate = FusedLeakyReLU(out_channel * sss) + + def forward(self, input, style, noise=None): + out = self.conv(input, style) + out = self.noise(out, noise=noise) + # out = out + self.bias + out = self.activate(out) + + return out + + +class ToRGB(nn.Module): + + def __init__(self, + in_channel, + style_dim, + upsample=True, + blur_kernel=[1, 3, 3, 1]): + super().__init__() + + if upsample: + self.upsample = Upsample(blur_kernel) + + self.conv = ModulatedConv2d( + in_channel, 3, 1, style_dim, demodulate=False) + self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1)) + + def forward(self, input, style, skip=None): + out = self.conv(input, style) + out = out + self.bias + + if skip is not None: + skip = self.upsample(skip) + + out = out + skip + + return out + + +class Generator(nn.Module): + + def __init__( + self, + size, + style_dim, + n_mlp, + channel_multiplier=2, + blur_kernel=[1, 3, 3, 1], + lr_mlp=0.01, + ): + super().__init__() + + self.size = size + self.n_mlp = n_mlp + self.style_dim = style_dim + + layers = [PixelNorm()] + + for i in range(n_mlp): + layers.append( + EqualLinear( + style_dim, + style_dim, + lr_mul=lr_mlp, + activation='fused_lrelu')) + + self.style = nn.Sequential(*layers) + + self.channels = { + 4: 512 // ratio, + 8: 512 // ratio, + 16: 512 // ratio, + 32: 512 // ratio, + 64: 256 // ratio * channel_multiplier, + 128: 128 // ratio * channel_multiplier, + 256: 64 // ratio * channel_multiplier, + 512: 32 // ratio * channel_multiplier, + 1024: 16 // ratio * channel_multiplier, + } + + self.input = ConstantInput(self.channels[4]) + self.conv1 = StyledConv( + self.channels[4], + self.channels[4], + 3, + style_dim, + blur_kernel=blur_kernel) + self.to_rgb1 = ToRGB(self.channels[4] * sss, style_dim, upsample=False) + + self.log_size = int(math.log(size, 2)) + + self.convs = nn.ModuleList() + self.upsamples = nn.ModuleList() + self.to_rgbs = nn.ModuleList() + + in_channel = self.channels[4] + + for i in range(3, self.log_size + 1): + out_channel = self.channels[2**i] + + self.convs.append( + StyledConv( + in_channel * sss, + out_channel, + 3, + style_dim, + upsample=True, + blur_kernel=blur_kernel, + )) + + self.convs.append( + StyledConv( + out_channel * sss, + out_channel, + 3, + style_dim, + blur_kernel=blur_kernel)) + + self.to_rgbs.append(ToRGB(out_channel * sss, style_dim)) + + in_channel = out_channel + + self.n_latent = self.log_size * 2 - 2 + + def make_noise(self): + device = self.input.input.device + + noises = [torch.randn(1, 1, 2**2, 2**2, device=device)] + + for i in range(3, self.log_size + 1): + for _ in range(2): + noises.append(torch.randn(1, 1, 2**i, 2**i, device=device)) + + return noises + + def mean_latent(self, n_latent): + latent_in = torch.randn( + n_latent, self.style_dim, device=self.input.input.device) + latent = self.style(latent_in).mean(0, keepdim=True) + + return latent + + def get_latent(self, input): + return self.style(input) + + def forward( + self, + styles, + return_latents=False, + inject_index=None, + truncation=1, + truncation_latent=None, + input_is_latent=False, + noise=None, + ): + if not input_is_latent: + styles = [self.style(s) for s in styles] + + if noise is None: + noise = [] + batch = styles[0].shape[0] + for i in range(self.n_mlp + 1): + size = 2**(i + 2) + noise.append( + torch.randn( + batch, + self.channels[size], + size, + size, + device=styles[0].device)) + + if truncation < 1: + style_t = [] + + for style in styles: + style_t.append(truncation_latent + + truncation * (style - truncation_latent)) + + styles = style_t + + if len(styles) < 2: + inject_index = self.n_latent + + latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) + + else: + if inject_index is None: + inject_index = random.randint(1, self.n_latent - 1) + + latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) + latent2 = styles[1].unsqueeze(1).repeat( + 1, self.n_latent - inject_index, 1) + + latent = torch.cat([latent, latent2], 1) + + out = self.input(latent) + out = self.conv1(out, latent[:, 0], noise=noise[0]) + + skip = self.to_rgb1(out, latent[:, 1]) + + i = 1 + noise_i = 1 + + for conv1, conv2, to_rgb in zip(self.convs[::2], self.convs[1::2], + self.to_rgbs): + out = conv1(out, latent[:, i], noise=noise[(noise_i + 1) // 2]) + out = conv2(out, latent[:, i + 1], noise=noise[(noise_i + 2) // 2]) + skip = to_rgb(out, latent[:, i + 2], skip) + + i += 2 + noise_i += 2 + + image = skip + + if return_latents: + return image, latent + + else: + return image, None + + +class ConvLayer(nn.Sequential): + + def __init__( + self, + in_channel, + out_channel, + kernel_size, + downsample=False, + blur_kernel=[1, 3, 3, 1], + bias=True, + activate=True, + ): + layers = [] + + if downsample: + factor = 2 + p = (len(blur_kernel) - factor) + (kernel_size - 1) + pad0 = (p + 1) // 2 + pad1 = p // 2 + + layers.append(Blur(blur_kernel, pad=(pad0, pad1))) + + stride = 2 + self.padding = 0 + + else: + stride = 1 + self.padding = kernel_size // 2 + + layers.append( + EqualConv2d( + in_channel, + out_channel, + kernel_size, + padding=self.padding, + stride=stride, + bias=bias and not activate, + )) + + if activate: + if bias: + layers.append(FusedLeakyReLU(out_channel)) + + else: + layers.append(ScaledLeakyReLU(0.2)) + + super().__init__(*layers) + + +class ResBlock(nn.Module): + + def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + self.conv1 = ConvLayer(in_channel, in_channel, 3) + self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True) + + self.skip = ConvLayer( + in_channel, + out_channel, + 1, + downsample=True, + activate=False, + bias=False) + + def forward(self, input): + out = self.conv1(input) + out = self.conv2(out) + + skip = self.skip(input) + out = (out + skip) / math.sqrt(2) + + return out + + +class Discriminator(nn.Module): + + def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]): + super().__init__() + + channels = { + 4: 512, + 8: 512, + 16: 512, + 32: 512, + 64: 256 * channel_multiplier, + 128: 128 * channel_multiplier, + 256: 64 * channel_multiplier, + 512: 32 * channel_multiplier, + 1024: 16 * channel_multiplier, + } + + convs = [ConvLayer(3, channels[size], 1)] + + log_size = int(math.log(size, 2)) + + in_channel = channels[size] + + for i in range(log_size, 2, -1): + out_channel = channels[2**(i - 1)] + + convs.append(ResBlock(in_channel, out_channel, blur_kernel)) + + in_channel = out_channel + + self.convs = nn.Sequential(*convs) + + self.stddev_group = 4 + self.stddev_feat = 1 + + self.final_conv = ConvLayer(in_channel + 1, channels[4], 3) + self.final_linear = nn.Sequential( + EqualLinear( + channels[4] * 4 * 4, channels[4], activation='fused_lrelu'), + EqualLinear(channels[4], 1), + ) + + def forward(self, input): + out = self.convs(input) + + batch, channel, height, width = out.shape + group = min(batch, self.stddev_group) + stddev = out.view(group, -1, self.stddev_feat, + channel // self.stddev_feat, height, width) + stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8) + stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2) + stddev = stddev.repeat(group, 1, height, width) + out = torch.cat([out, stddev], 1) + + out = self.final_conv(out) + + out = out.view(batch, -1) + out = self.final_linear(out) + return out + + +class FullGenerator(nn.Module): + + def __init__( + self, + size, + style_dim, + n_mlp, + channel_multiplier=2, + blur_kernel=[1, 3, 3, 1], + lr_mlp=0.01, + ): + super().__init__() + channels = { + 4: 512 // ratio, + 8: 512 // ratio, + 16: 512 // ratio, + 32: 512 // ratio, + 64: 256 // ratio * channel_multiplier, + 128: 128 // ratio * channel_multiplier, + 256: 64 // ratio * channel_multiplier, + 512: 32 // ratio * channel_multiplier, + 1024: 16 // ratio * channel_multiplier, + } + + self.log_size = int(math.log(size, 2)) + self.generator = Generator( + size, + style_dim, + n_mlp, + channel_multiplier=channel_multiplier, + blur_kernel=blur_kernel, + lr_mlp=lr_mlp) + + conv = [ConvLayer(3, channels[size], 1)] + self.ecd0 = nn.Sequential(*conv) + in_channel = channels[size] + + self.names = ['ecd%d' % i for i in range(self.log_size - 1)] + for i in range(self.log_size, 2, -1): + out_channel = channels[2**(i - 1)] + conv = [ConvLayer(in_channel, out_channel, 3, downsample=True)] + setattr(self, self.names[self.log_size - i + 1], + nn.Sequential(*conv)) + in_channel = out_channel + self.final_linear = nn.Sequential( + EqualLinear( + channels[4] * 4 * 4, style_dim, activation='fused_lrelu')) + + def forward( + self, + inputs, + return_latents=False, + inject_index=None, + truncation=1, + truncation_latent=None, + input_is_latent=False, + ): + noise = [] + for i in range(self.log_size - 1): + ecd = getattr(self, self.names[i]) + inputs = ecd(inputs) + noise.append(inputs) + inputs = inputs.view(inputs.shape[0], -1) + outs = self.final_linear(inputs) + outs = self.generator([outs], + return_latents, + inject_index, + truncation, + truncation_latent, + input_is_latent, + noise=noise[::-1]) + + return outs diff --git a/modelscope/models/cv/image_face_fusion/facegan/op/__init__.py b/modelscope/models/cv/image_face_fusion/facegan/op/__init__.py new file mode 100644 index 00000000..74477cfb --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/facegan/op/__init__.py @@ -0,0 +1,4 @@ +# The implementation is adopted from stylegan2-pytorch, made public available under the MIT License +# at https://github.com/rosinality/stylegan2-pytorch +from .fused_act import FusedLeakyReLU, fused_leaky_relu +from .upfirdn2d import upfirdn2d diff --git a/modelscope/models/cv/image_face_fusion/facegan/op/conv2d_gradfix.py b/modelscope/models/cv/image_face_fusion/facegan/op/conv2d_gradfix.py new file mode 100644 index 00000000..a3aba91f --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/facegan/op/conv2d_gradfix.py @@ -0,0 +1,228 @@ +# The implementation is adopted from stylegan2-pytorch, made public available under the MIT License +# at https://github.com/rosinality/stylegan2-pytorch/blob/master/op/conv2d_gradfix.py +import contextlib +import warnings + +import torch +from torch import autograd +from torch.nn import functional as F + +enabled = True +weight_gradients_disabled = False + + +@contextlib.contextmanager +def no_weight_gradients(): + global weight_gradients_disabled + + old = weight_gradients_disabled + weight_gradients_disabled = True + yield + weight_gradients_disabled = old + + +def conv2d(input, + weight, + bias=None, + stride=1, + padding=0, + dilation=1, + groups=1): + if could_use_op(input): + return conv2d_gradfix( + transpose=False, + weight_shape=weight.shape, + stride=stride, + padding=padding, + output_padding=0, + dilation=dilation, + groups=groups, + ).apply(input, weight, bias) + + return F.conv2d( + input=input, + weight=weight, + bias=bias, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + ) + + +def conv_transpose2d( + input, + weight, + bias=None, + stride=1, + padding=0, + output_padding=0, + groups=1, + dilation=1, +): + if could_use_op(input): + return conv2d_gradfix( + transpose=True, + weight_shape=weight.shape, + stride=stride, + padding=padding, + output_padding=output_padding, + groups=groups, + dilation=dilation, + ).apply(input, weight, bias) + + return F.conv_transpose2d( + input=input, + weight=weight, + bias=bias, + stride=stride, + padding=padding, + output_padding=output_padding, + dilation=dilation, + groups=groups, + ) + + +def could_use_op(input): + if (not enabled) or (not torch.backends.cudnn.enabled): + return False + + if input.device.type != 'cuda': + return False + + warnings.warn( + f'conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d().' + ) + + return False + + +def ensure_tuple(xs, ndim): + xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs, ) * ndim + + return xs + + +conv2d_gradfix_cache = dict() + + +def conv2d_gradfix(transpose, weight_shape, stride, padding, output_padding, + dilation, groups): + ndim = 2 + weight_shape = tuple(weight_shape) + stride = ensure_tuple(stride, ndim) + padding = ensure_tuple(padding, ndim) + output_padding = ensure_tuple(output_padding, ndim) + dilation = ensure_tuple(dilation, ndim) + + key = (transpose, weight_shape, stride, padding, output_padding, dilation, + groups) + if key in conv2d_gradfix_cache: + return conv2d_gradfix_cache[key] + + common_kwargs = dict( + stride=stride, padding=padding, dilation=dilation, groups=groups) + + def calc_output_padding(input_shape, output_shape): + if transpose: + return [0, 0] + + a = input_shape[i + 2] - (output_shape[i + 2] - 1) * stride[i] + return [ + a - (1 - 2 * padding[i]) - dilation[i] * (weight_shape[i + 2] - 1) + for i in range(ndim) + ] + + class Conv2d(autograd.Function): + + @staticmethod + def forward(ctx, input, weight, bias): + if not transpose: + out = F.conv2d( + input=input, weight=weight, bias=bias, **common_kwargs) + + else: + out = F.conv_transpose2d( + input=input, + weight=weight, + bias=bias, + output_padding=output_padding, + **common_kwargs, + ) + + ctx.save_for_backward(input, weight) + + return out + + @staticmethod + def backward(ctx, grad_output): + input, weight = ctx.saved_tensors + grad_input, grad_weight, grad_bias = None, None, None + + if ctx.needs_input_grad[0]: + p = calc_output_padding( + input_shape=input.shape, output_shape=grad_output.shape) + grad_input = conv2d_gradfix( + transpose=(not transpose), + weight_shape=weight_shape, + output_padding=p, + **common_kwargs, + ).apply(grad_output, weight, None) + + if ctx.needs_input_grad[1] and not weight_gradients_disabled: + grad_weight = Conv2dGradWeight.apply(grad_output, input) + + if ctx.needs_input_grad[2]: + grad_bias = grad_output.sum((0, 2, 3)) + + return grad_input, grad_weight, grad_bias + + class Conv2dGradWeight(autograd.Function): + + @staticmethod + def forward(ctx, grad_output, input): + op = torch._C._jit_get_operation( + 'aten::cudnn_convolution_backward_weight' if not transpose else + 'aten::cudnn_convolution_transpose_backward_weight') + flags = [ + torch.backends.cudnn.benchmark, + torch.backends.cudnn.deterministic, + torch.backends.cudnn.allow_tf32, + ] + grad_weight = op( + weight_shape, + grad_output, + input, + padding, + stride, + dilation, + groups, + *flags, + ) + ctx.save_for_backward(grad_output, input) + + return grad_weight + + @staticmethod + def backward(ctx, grad_grad_weight): + grad_output, input = ctx.saved_tensors + grad_grad_output, grad_grad_input = None, None + + if ctx.needs_input_grad[0]: + grad_grad_output = Conv2d.apply(input, grad_grad_weight, None) + + if ctx.needs_input_grad[1]: + p = calc_output_padding( + input_shape=input.shape, output_shape=grad_output.shape) + grad_grad_input = conv2d_gradfix( + transpose=(not transpose), + weight_shape=weight_shape, + output_padding=p, + **common_kwargs, + ).apply(grad_output, grad_grad_weight, None) + + return grad_grad_output, grad_grad_input + + conv2d_gradfix_cache[key] = Conv2d + + return Conv2d diff --git a/modelscope/models/cv/image_face_fusion/facegan/op/fused_act.py b/modelscope/models/cv/image_face_fusion/facegan/op/fused_act.py new file mode 100644 index 00000000..7db45d3c --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/facegan/op/fused_act.py @@ -0,0 +1,113 @@ +# The implementation is adopted from stylegan2-pytorch, made public available under the MIT License +# at https://github.com/rosinality/stylegan2-pytorch/blob/master/op/fused_act.py +import torch +from torch import nn +from torch.autograd import Function +from torch.nn import functional as F + +def_lib = False + + +class FusedLeakyReLUFunctionBackward(Function): + + @staticmethod + def forward(ctx, grad_output, out, bias, negative_slope, scale): + ctx.save_for_backward(out) + ctx.negative_slope = negative_slope + ctx.scale = scale + + empty = grad_output.new_empty(0) + + grad_input = fused.fused_bias_act(grad_output.contiguous(), empty, out, + 3, 1, negative_slope, scale) + + dim = [0] + + if grad_input.ndim > 2: + dim += list(range(2, grad_input.ndim)) + + if bias: + grad_bias = grad_input.sum(dim).detach() + + else: + grad_bias = empty + + return grad_input, grad_bias + + @staticmethod + def backward(ctx, gradgrad_input, gradgrad_bias): + out, = ctx.saved_tensors + gradgrad_out = fused.fused_bias_act( + gradgrad_input.contiguous(), + gradgrad_bias, + out, + 3, + 1, + ctx.negative_slope, + ctx.scale, + ) + + return gradgrad_out, None, None, None, None + + +class FusedLeakyReLUFunction(Function): + + @staticmethod + def forward(ctx, input, bias, negative_slope, scale): + empty = input.new_empty(0) + + ctx.bias = bias is not None + + if bias is None: + bias = empty + + out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, + scale) + ctx.save_for_backward(out) + ctx.negative_slope = negative_slope + ctx.scale = scale + + return out + + @staticmethod + def backward(ctx, grad_output): + out, = ctx.saved_tensors + + grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply( + grad_output, out, ctx.bias, ctx.negative_slope, ctx.scale) + + if not ctx.bias: + grad_bias = None + + return grad_input, grad_bias, None, None + + +class FusedLeakyReLU(nn.Module): + + def __init__(self, channel, bias=True, negative_slope=0.2, scale=2**0.5): + super().__init__() + + if bias: + self.bias = nn.Parameter(torch.zeros(channel)) + + else: + self.bias = None + + self.negative_slope = negative_slope + self.scale = scale + + def forward(self, input): + return fused_leaky_relu(input, self.bias, self.negative_slope, + self.scale) + + +def fused_leaky_relu(input, bias=None, negative_slope=0.2, scale=2**0.5): + if not def_lib: + if bias is not None: + rest_dim = [1] * (input.ndim - bias.ndim - 1) + return (F.leaky_relu( + input + bias.view(1, bias.shape[0], *rest_dim), + negative_slope=0.2) * scale) + + else: + return F.leaky_relu(input, negative_slope=0.2) * scale diff --git a/modelscope/models/cv/image_face_fusion/facegan/op/upfirdn2d.py b/modelscope/models/cv/image_face_fusion/facegan/op/upfirdn2d.py new file mode 100644 index 00000000..3b93c082 --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/facegan/op/upfirdn2d.py @@ -0,0 +1,198 @@ +# The implementation is adopted from stylegan2-pytorch, made public available under the MIT License +# at https://github.com/rosinality/stylegan2-pytorch/blob/master/op/upfirdn2d.py +from collections import abc + +import torch +from torch.autograd import Function +from torch.nn import functional as F + +def_lib = False + + +class UpFirDn2dBackward(Function): + + @staticmethod + def forward(ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, + in_size, out_size): + + up_x, up_y = up + down_x, down_y = down + g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad + + grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1) + + grad_input = upfirdn2d_op.upfirdn2d( + grad_output, + grad_kernel, + down_x, + down_y, + up_x, + up_y, + g_pad_x0, + g_pad_x1, + g_pad_y0, + g_pad_y1, + ) + grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], + in_size[3]) + + ctx.save_for_backward(kernel) + + pad_x0, pad_x1, pad_y0, pad_y1 = pad + + ctx.up_x = up_x + ctx.up_y = up_y + ctx.down_x = down_x + ctx.down_y = down_y + ctx.pad_x0 = pad_x0 + ctx.pad_x1 = pad_x1 + ctx.pad_y0 = pad_y0 + ctx.pad_y1 = pad_y1 + ctx.in_size = in_size + ctx.out_size = out_size + + return grad_input + + @staticmethod + def backward(ctx, gradgrad_input): + kernel, = ctx.saved_tensors + + gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], + ctx.in_size[3], 1) + + gradgrad_out = upfirdn2d_op.upfirdn2d( + gradgrad_input, + kernel, + ctx.up_x, + ctx.up_y, + ctx.down_x, + ctx.down_y, + ctx.pad_x0, + ctx.pad_x1, + ctx.pad_y0, + ctx.pad_y1, + ) + # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3]) + gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.in_size[1], + ctx.out_size[0], ctx.out_size[1]) + + return gradgrad_out, None, None, None, None, None, None, None, None + + +class UpFirDn2d(Function): + + @staticmethod + def forward(ctx, input, kernel, up, down, pad): + up_x, up_y = up + down_x, down_y = down + pad_x0, pad_x1, pad_y0, pad_y1 = pad + + kernel_h, kernel_w = kernel.shape + batch, channel, in_h, in_w = input.shape + ctx.in_size = input.shape + + input = input.reshape(-1, in_h, in_w, 1) + + ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1])) + + out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h + down_y) // down_y + out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w + down_x) // down_x + ctx.out_size = (out_h, out_w) + + ctx.up = (up_x, up_y) + ctx.down = (down_x, down_y) + ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1) + + g_pad_x0 = kernel_w - pad_x0 - 1 + g_pad_y0 = kernel_h - pad_y0 - 1 + g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1 + g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1 + + ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1) + + out = upfirdn2d_op.upfirdn2d(input, kernel, up_x, up_y, down_x, down_y, + pad_x0, pad_x1, pad_y0, pad_y1) + # out = out.view(major, out_h, out_w, minor) + out = out.view(-1, channel, out_h, out_w) + + return out + + @staticmethod + def backward(ctx, grad_output): + kernel, grad_kernel = ctx.saved_tensors + + grad_input = None + + if ctx.needs_input_grad[0]: + grad_input = UpFirDn2dBackward.apply( + grad_output, + kernel, + grad_kernel, + ctx.up, + ctx.down, + ctx.pad, + ctx.g_pad, + ctx.in_size, + ctx.out_size, + ) + + return grad_input, None, None, None, None + + +def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)): + if not isinstance(up, abc.Iterable): + up = (up, up) + + if not isinstance(down, abc.Iterable): + down = (down, down) + + if len(pad) == 2: + pad = (pad[0], pad[1], pad[0], pad[1]) + + if not def_lib: + out = upfirdn2d_native(input, kernel, *up, *down, *pad) + + return out + + +def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, + pad_y0, pad_y1): + _, channel, in_h, in_w = input.shape + input = input.reshape(-1, in_h, in_w, 1) + + _, in_h, in_w, minor = input.shape + kernel_h, kernel_w = kernel.shape + + out = input.view(-1, in_h, 1, in_w, 1, minor) + out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1]) + out = out.view(-1, in_h * up_y, in_w * up_x, minor) + + out = F.pad( + out, + [0, 0, + max(pad_x0, 0), + max(pad_x1, 0), + max(pad_y0, 0), + max(pad_y1, 0)]) + out = out[:, + max(-pad_y0, 0):out.shape[1] - max(-pad_y1, 0), + max(-pad_x0, 0):out.shape[2] - max(-pad_x1, 0)] + + out = out.permute(0, 3, 1, 2) + out = out.reshape( + [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]) + w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w) + out = F.conv2d(out, w) + out = out.reshape( + -1, + minor, + in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1, + in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, + ) + out = out.permute(0, 2, 3, 1) + out = out[:, ::down_y, ::down_x, :] + + out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h + down_y) // down_y + out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w + down_x) // down_x + + return out.view(-1, channel, out_h, out_w) diff --git a/modelscope/models/cv/image_face_fusion/facelib/__init__.py b/modelscope/models/cv/image_face_fusion/facelib/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/image_face_fusion/facelib/align_trans.py b/modelscope/models/cv/image_face_fusion/facelib/align_trans.py new file mode 100644 index 00000000..554b0e7c --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/facelib/align_trans.py @@ -0,0 +1,301 @@ +# The implementation here is modified based on InsightFace_Pytorch, originally MIT License and publicly available +# at https://github.com/TreB1eN/InsightFace_Pytorch/blob/master/mtcnn_pytorch/src/align_trans.py +import cv2 +import numpy as np + +from .matlab_cp2tform import get_similarity_transform_for_cv2 + +# reference facial points, a list of coordinates (x,y) +REFERENCE_FACIAL_POINTS = [[30.29459953, 51.69630051], + [65.53179932, 51.50139999], + [48.02519989, + 71.73660278], [33.54930115, 92.3655014], + [62.72990036, 92.20410156]] + +DEFAULT_CROP_SIZE = (96, 112) + + +class FaceWarpException(Exception): + + def __str__(self): + return 'In File {}:{}'.format(__file__, super.__str__(self)) + + +def get_reference_facial_points(output_size=None, + inner_padding_factor=0.0, + outer_padding=(0, 0), + default_square=False): + """ + Function: + ---------- + get reference 5 key points according to crop settings: + 0. Set default crop_size: + if default_square: + crop_size = (112, 112) + else: + crop_size = (96, 112) + 1. Pad the crop_size by inner_padding_factor in each side; + 2. Resize crop_size into (output_size - outer_padding*2), + pad into output_size with outer_padding; + 3. Output reference_5point; + Parameters: + ---------- + @output_size: (w, h) or None + size of aligned face image + @inner_padding_factor: (w_factor, h_factor) + padding factor for inner (w, h) + @outer_padding: (w_pad, h_pad) + each row is a pair of coordinates (x, y) + @default_square: True or False + if True: + default crop_size = (112, 112) + else: + default crop_size = (96, 112); + !!! make sure, if output_size is not None: + (output_size - outer_padding) + = some_scale * (default crop_size * (1.0 + inner_padding_factor)) + Returns: + ---------- + @reference_5point: 5x2 np.array + each row is a pair of transformed coordinates (x, y) + """ + + tmp_5pts = np.array(REFERENCE_FACIAL_POINTS) + tmp_crop_size = np.array(DEFAULT_CROP_SIZE) + + # 0) make the inner region a square + if default_square: + size_diff = max(tmp_crop_size) - tmp_crop_size + tmp_5pts += size_diff / 2 + tmp_crop_size += size_diff + + if (output_size and output_size[0] == tmp_crop_size[0] + and output_size[1] == tmp_crop_size[1]): + return tmp_5pts + + if (inner_padding_factor == 0 and outer_padding == (0, 0)): + if output_size is None: + return tmp_5pts + else: + raise FaceWarpException( + 'No paddings to do, output_size must be None or {}'.format( + tmp_crop_size)) + + if not (0 <= inner_padding_factor <= 1.0): + raise FaceWarpException('Not (0 <= inner_padding_factor <= 1.0)') + + if ((inner_padding_factor > 0 or outer_padding[0] > 0 + or outer_padding[1] > 0) and output_size is None): + output_size = tmp_crop_size * \ + (1 + inner_padding_factor * 2).astype(np.int32) + output_size += np.array(outer_padding) + + if not (outer_padding[0] < output_size[0] + and outer_padding[1] < output_size[1]): + raise FaceWarpException('Not (outer_padding[0] < output_size[0]' + 'and outer_padding[1] < output_size[1])') + + # 1) pad the inner region according inner_padding_factor + if inner_padding_factor > 0: + size_diff = tmp_crop_size * inner_padding_factor * 2 + tmp_5pts += size_diff / 2 + tmp_crop_size += np.round(size_diff).astype(np.int32) + + # 2) resize the padded inner region + size_bf_outer_pad = np.array(output_size) - np.array(outer_padding) * 2 + + if size_bf_outer_pad[0] * tmp_crop_size[1] != size_bf_outer_pad[ + 1] * tmp_crop_size[0]: + raise FaceWarpException( + 'Must have (output_size - outer_padding)' + '= some_scale * (crop_size * (1.0 + inner_padding_factor)') + + scale_factor = size_bf_outer_pad[0].astype(np.float32) / tmp_crop_size[0] + tmp_5pts = tmp_5pts * scale_factor + tmp_crop_size = size_bf_outer_pad + + # 3) add outer_padding to make output_size + reference_5point = tmp_5pts + np.array(outer_padding) + tmp_crop_size = output_size + + return reference_5point + + +def get_affine_transform_matrix(src_pts, dst_pts): + """ + Function: + ---------- + get affine transform matrix 'tfm' from src_pts to dst_pts + Parameters: + ---------- + @src_pts: Kx2 np.array + source points matrix, each row is a pair of coordinates (x, y) + @dst_pts: Kx2 np.array + destination points matrix, each row is a pair of coordinates (x, y) + Returns: + ---------- + @tfm: 2x3 np.array + transform matrix from src_pts to dst_pts + """ + + tfm = np.float32([[1, 0, 0], [0, 1, 0]]) + n_pts = src_pts.shape[0] + ones = np.ones((n_pts, 1), src_pts.dtype) + src_pts_ = np.hstack([src_pts, ones]) + dst_pts_ = np.hstack([dst_pts, ones]) + + A, res, rank, s = np.linalg.lstsq(src_pts_, dst_pts_) + + if rank == 3: + tfm = np.float32([[A[0, 0], A[1, 0], A[2, 0]], + [A[0, 1], A[1, 1], A[2, 1]]]) + elif rank == 2: + tfm = np.float32([[A[0, 0], A[1, 0], 0], [A[0, 1], A[1, 1], 0]]) + + return tfm + + +def warp_and_crop_face(src_img, + facial_pts, + reference_pts=None, + crop_size=(96, 112), + align_type='smilarity', + return_trans_inv=False): + """ + Function: + ---------- + apply affine transform 'trans' to uv + Parameters: + ---------- + @src_img: 3x3 np.array + input image + @facial_pts: could be + 1)a list of K coordinates (x,y) + or + 2) Kx2 or 2xK np.array + each row or col is a pair of coordinates (x, y) + @reference_pts: could be + 1) a list of K coordinates (x,y) + or + 2) Kx2 or 2xK np.array + each row or col is a pair of coordinates (x, y) + or + 3) None + if None, use default reference facial points + @crop_size: (w, h) + output face image size + @align_type: transform type, could be one of + 1) 'similarity': use similarity transform + 2) 'cv2_affine': use the first 3 points to do affine transform, + by calling cv2.getAffineTransform() + 3) 'affine': use all points to do affine transform + Returns: + ---------- + @face_img: output face image with size (w, h) = @crop_size + """ + + if reference_pts is None: + if crop_size[0] == 96 and crop_size[1] == 112: + reference_pts = REFERENCE_FACIAL_POINTS + else: + default_square = False + inner_padding_factor = 0 + outer_padding = (0, 0) + output_size = crop_size + + reference_pts = get_reference_facial_points( + output_size, inner_padding_factor, outer_padding, + default_square) + + ref_pts = np.float32(reference_pts) + ref_pts = (ref_pts - 112 / 2) * 0.85 + 112 / 2 + ref_pts *= crop_size[0] / 112. + ref_pts_shp = ref_pts.shape + if max(ref_pts_shp) < 3 or min(ref_pts_shp) != 2: + raise FaceWarpException( + 'reference_pts.shape must be (K,2) or (2,K) and K>2') + + if ref_pts_shp[0] == 2: + ref_pts = ref_pts.T + + src_pts = np.float32(facial_pts) + src_pts_shp = src_pts.shape + if max(src_pts_shp) < 3 or min(src_pts_shp) != 2: + raise FaceWarpException( + 'facial_pts.shape must be (K,2) or (2,K) and K>2') + + if src_pts_shp[0] == 2: + src_pts = src_pts.T + + if src_pts.shape != ref_pts.shape: + raise FaceWarpException( + 'facial_pts and reference_pts must have the same shape') + + if align_type == 'cv2_affine': + tfm = cv2.getAffineTransform(src_pts[0:3], ref_pts[0:3]) + elif align_type == 'affine': + tfm = get_affine_transform_matrix(src_pts, ref_pts) + else: + tfm, tfm_inv = get_similarity_transform_for_cv2(src_pts, ref_pts) + + face_img = cv2.warpAffine(src_img, tfm, (crop_size[0], crop_size[1])) + + if return_trans_inv: + return face_img, tfm_inv + else: + return face_img + + +def get_f5p(landmarks, np_img): + eye_left = find_pupil(landmarks[36:41], np_img) + eye_right = find_pupil(landmarks[42:47], np_img) + if eye_left is None or eye_right is None: + print('cannot find 5 points with find_pupil, used mean instead.!') + eye_left = landmarks[36:41].mean(axis=0) + eye_right = landmarks[42:47].mean(axis=0) + nose = landmarks[30] + mouth_left = landmarks[48] + mouth_right = landmarks[54] + f5p = [[eye_left[0], eye_left[1]], [eye_right[0], eye_right[1]], + [nose[0], nose[1]], [mouth_left[0], mouth_left[1]], + [mouth_right[0], mouth_right[1]]] + return f5p + + +def find_pupil(landmarks, np_img): + h, w, _ = np_img.shape + xmax = int(landmarks[:, 0].max()) + xmin = int(landmarks[:, 0].min()) + ymax = int(landmarks[:, 1].max()) + ymin = int(landmarks[:, 1].min()) + + if ymin >= ymax or xmin >= xmax or ymin < 0 or xmin < 0 or ymax > h or xmax > w: + return None + eye_img_bgr = np_img[ymin:ymax, xmin:xmax, :] + eye_img = cv2.cvtColor(eye_img_bgr, cv2.COLOR_BGR2GRAY) + eye_img = cv2.equalizeHist(eye_img) + n_marks = landmarks - np.array([xmin, ymin]).reshape([1, 2]) + eye_mask = cv2.fillConvexPoly( + np.zeros_like(eye_img), n_marks.astype(np.int32), 1) + ret, thresh = cv2.threshold(eye_img, 100, 255, + cv2.THRESH_BINARY | cv2.THRESH_OTSU) + thresh = (1 - thresh / 255.) * eye_mask + cnt = 0 + xm = [] + ym = [] + for i in range(thresh.shape[0]): + for j in range(thresh.shape[1]): + if thresh[i, j] > 0.5: + xm.append(j) + ym.append(i) + cnt += 1 + if cnt != 0: + xm.sort() + ym.sort() + xm = xm[cnt // 2] + ym = ym[cnt // 2] + else: + xm = thresh.shape[1] / 2 + ym = thresh.shape[0] / 2 + + return xm + xmin, ym + ymin diff --git a/modelscope/models/cv/image_face_fusion/facelib/matlab_cp2tform.py b/modelscope/models/cv/image_face_fusion/facelib/matlab_cp2tform.py new file mode 100644 index 00000000..87b9fd1e --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/facelib/matlab_cp2tform.py @@ -0,0 +1,230 @@ +# The implementation is adopted from InsightFace_Pytorch, made publicly available under the MIT License +# at https://github.com/TreB1eN/InsightFace_Pytorch/blob/master/mtcnn_pytorch/src/matlab_cp2tform.py + +import numpy as np +from numpy.linalg import inv, lstsq +from numpy.linalg import matrix_rank as rank +from numpy.linalg import norm + + +class MatlabCp2tormException(Exception): + + def __str__(self): + return 'In File {}:{}'.format(__file__, super.__str__(self)) + + +def tformfwd(trans, uv): + """ + Function: + ---------- + apply affine transform 'trans' to uv + + Parameters: + ---------- + @trans: 3x3 np.array + transform matrix + @uv: Kx2 np.array + each row is a pair of coordinates (x, y) + + Returns: + ---------- + @xy: Kx2 np.array + each row is a pair of transformed coordinates (x, y) + """ + uv = np.hstack((uv, np.ones((uv.shape[0], 1)))) + xy = np.dot(uv, trans) + xy = xy[:, 0:-1] + return xy + + +def tforminv(trans, uv): + """ + Function: + ---------- + apply the inverse of affine transform 'trans' to uv + + Parameters: + ---------- + @trans: 3x3 np.array + transform matrix + @uv: Kx2 np.array + each row is a pair of coordinates (x, y) + + Returns: + ---------- + @xy: Kx2 np.array + each row is a pair of inverse-transformed coordinates (x, y) + """ + Tinv = inv(trans) + xy = tformfwd(Tinv, uv) + return xy + + +def findNonreflectiveSimilarity(uv, xy, options=None): + + options = {'K': 2} + + K = options['K'] + M = xy.shape[0] + x = xy[:, 0].reshape((-1, 1)) + y = xy[:, 1].reshape((-1, 1)) + + tmp1 = np.hstack((x, y, np.ones((M, 1)), np.zeros((M, 1)))) + tmp2 = np.hstack((y, -x, np.zeros((M, 1)), np.ones((M, 1)))) + X = np.vstack((tmp1, tmp2)) + + u = uv[:, 0].reshape((-1, 1)) + v = uv[:, 1].reshape((-1, 1)) + U = np.vstack((u, v)) + + if rank(X) >= 2 * K: + r, _, _, _ = lstsq(X, U) + r = np.squeeze(r) + else: + raise Exception('cp2tform:twoUniquePointsReq') + + sc = r[0] + ss = r[1] + tx = r[2] + ty = r[3] + + Tinv = np.array([[sc, -ss, 0], [ss, sc, 0], [tx, ty, 1]]) + + T = inv(Tinv) + T[:, 2] = np.array([0, 0, 1]) + + return T, Tinv + + +def findSimilarity(uv, xy, options=None): + + options = {'K': 2} + + trans1, trans1_inv = findNonreflectiveSimilarity(uv, xy, options) + + xyR = xy + xyR[:, 0] = -1 * xyR[:, 0] + + trans2r, trans2r_inv = findNonreflectiveSimilarity(uv, xyR, options) + + TreflectY = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]]) + + trans2 = np.dot(trans2r, TreflectY) + + xy1 = tformfwd(trans1, uv) + norm1 = norm(xy1 - xy) + + xy2 = tformfwd(trans2, uv) + norm2 = norm(xy2 - xy) + + if norm1 <= norm2: + return trans1, trans1_inv + else: + trans2_inv = inv(trans2) + return trans2, trans2_inv + + +def get_similarity_transform(src_pts, dst_pts, reflective=True): + """ + Function: + ---------- + Find Similarity Transform Matrix 'trans': + u = src_pts[:, 0] + v = src_pts[:, 1] + x = dst_pts[:, 0] + y = dst_pts[:, 1] + [x, y, 1] = [u, v, 1] * trans + + Parameters: + ---------- + @src_pts: Kx2 np.array + source points, each row is a pair of coordinates (x, y) + @dst_pts: Kx2 np.array + destination points, each row is a pair of transformed + coordinates (x, y) + @reflective: True or False + if True: + use reflective similarity transform + else: + use non-reflective similarity transform + + Returns: + ---------- + @trans: 3x3 np.array + transform matrix from uv to xy + trans_inv: 3x3 np.array + inverse of trans, transform matrix from xy to uv + """ + + if reflective: + trans, trans_inv = findSimilarity(src_pts, dst_pts) + else: + trans, trans_inv = findNonreflectiveSimilarity(src_pts, dst_pts) + + return trans, trans_inv + + +def cvt_tform_mat_for_cv2(trans): + """ + Function: + ---------- + Convert Transform Matrix 'trans' into 'cv2_trans' which could be + directly used by cv2.warpAffine(): + u = src_pts[:, 0] + v = src_pts[:, 1] + x = dst_pts[:, 0] + y = dst_pts[:, 1] + [x, y].T = cv_trans * [u, v, 1].T + + Parameters: + ---------- + @trans: 3x3 np.array + transform matrix from uv to xy + + Returns: + ---------- + @cv2_trans: 2x3 np.array + transform matrix from src_pts to dst_pts, could be directly used + for cv2.warpAffine() + """ + cv2_trans = trans[:, 0:2].T + + return cv2_trans + + +def get_similarity_transform_for_cv2(src_pts, dst_pts, reflective=True): + """ + Function: + ---------- + Find Similarity Transform Matrix 'cv2_trans' which could be + directly used by cv2.warpAffine(): + u = src_pts[:, 0] + v = src_pts[:, 1] + x = dst_pts[:, 0] + y = dst_pts[:, 1] + [x, y].T = cv_trans * [u, v, 1].T + + Parameters: + ---------- + @src_pts: Kx2 np.array + source points, each row is a pair of coordinates (x, y) + @dst_pts: Kx2 np.array + destination points, each row is a pair of transformed + coordinates (x, y) + reflective: True or False + if True: + use reflective similarity transform + else: + use non-reflective similarity transform + + Returns: + ---------- + @cv2_trans: 2x3 np.array + transform matrix from src_pts to dst_pts, could be directly used + for cv2.warpAffine() + """ + trans, trans_inv = get_similarity_transform(src_pts, dst_pts, reflective) + cv2_trans = cvt_tform_mat_for_cv2(trans) + cv2_trans_inv = cvt_tform_mat_for_cv2(trans_inv) + + return cv2_trans, cv2_trans_inv diff --git a/modelscope/models/cv/image_face_fusion/image_face_fusion.py b/modelscope/models/cv/image_face_fusion/image_face_fusion.py new file mode 100644 index 00000000..24907ceb --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/image_face_fusion.py @@ -0,0 +1,253 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os +from collections import OrderedDict +from typing import Any, Dict + +import cv2 +import numpy as np +import PIL.Image as Image +import torch +import torch.nn.functional as F +import torchvision.transforms as transforms + +from modelscope.metainfo import Models +from modelscope.models.base import Tensor, TorchModel +from modelscope.models.builder import MODELS +from modelscope.models.cv.face_detection.peppa_pig_face.facer import FaceAna +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.logger import get_logger +from .facegan.gan_wrap import GANWrap +from .facelib.align_trans import (get_f5p, get_reference_facial_points, + warp_and_crop_face) +from .network.aei_flow_net import AEI_Net +from .network.bfm import ParametricFaceModel +from .network.facerecon_model import ReconNetWrapper +from .network.model_irse import Backbone +from .network.ops import warp_affine_torch + +logger = get_logger() + +__all__ = ['ImageFaceFusion'] + + +@MODELS.register_module( + Tasks.image_face_fusion, module_name=Models.image_face_fusion) +class ImageFaceFusion(TorchModel): + + def __init__(self, model_dir: str, *args, **kwargs): + """initialize the image face fusion model from the `model_dir` path. + + Args: + model_dir (str): the model path. + """ + super().__init__(model_dir, *args, **kwargs) + + if torch.cuda.is_available(): + self.device = torch.device('cuda') + else: + self.device = torch.device('cpu') + + self.num_kp = 17 + self.id_dim = 512 + + self.netG = AEI_Net( + c_id=self.id_dim, num_kp=self.num_kp, device=self.device) + model_path = os.path.join(model_dir, ModelFile.TORCH_MODEL_FILE) + checkpoints = torch.load(model_path, map_location='cpu') + model_state = self.convert_state_dict(checkpoints['state_dict']) + self.netG.load_state_dict(model_state) + self.netG = self.netG.to(self.device) + self.netG.eval() + + self.arcface = Backbone([112, 112], 100, 'ir') + arcface_path = os.path.join(model_dir, 'faceRecog', + 'CurricularFace_Backbone.pth') + self.arcface.load_state_dict( + torch.load(arcface_path, map_location='cpu'), strict=False) + self.arcface = self.arcface.to(self.device) + self.arcface.eval() + + self.f_3d = ReconNetWrapper(net_recon='resnet50', use_last_fc=False) + f_3d_path = os.path.join(model_dir, '3dRecon', 'face_3d.pth') + self.f_3d.load_state_dict( + torch.load(f_3d_path, map_location='cpu')['net_recon']) + self.f_3d = self.f_3d.to(self.device) + self.f_3d.eval() + + bfm_dir = os.path.join(model_dir, 'BFM') + self.face_model = ParametricFaceModel(bfm_folder=bfm_dir) + self.face_model.to(self.device) + + face_enhance_path = os.path.join(model_dir, 'faceEnhance', + '350000-Ns256.pt') + self.ganwrap = GANWrap( + model_path=face_enhance_path, + size=256, + channel_multiplier=1, + device=self.device) + + self.facer = FaceAna(model_dir) + + logger.info('load facefusion models done') + + self.mask_init = cv2.imread(os.path.join(model_dir, 'alpha.jpg')) + self.mask_init = cv2.resize(self.mask_init, (256, 256)) + self.mask = self.image_transform(self.mask_init, is_norm=False) + + self.test_transform = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) + ]) + + logger.info('init done') + + def convert_state_dict(self, state_dict): + if not next(iter(state_dict)).startswith('module.'): + return state_dict + new_state_dict = OrderedDict() + + split_index = 0 + for cur_key, cur_value in state_dict.items(): + if cur_key.startswith('module.model'): + split_index = 13 + elif cur_key.startswith('module'): + split_index = 7 + + break + + for k, v in state_dict.items(): + name = k[split_index:] + new_state_dict[name] = v + return new_state_dict + + def image_transform(self, + image, + is_norm=True, + mean=(0.5, 0.5, 0.5), + std=(0.5, 0.5, 0.5)): + image = image.astype(np.float32) + image = image / 255.0 + if is_norm: + image -= mean + image /= std + + image = image.transpose((2, 0, 1)) + image = np.expand_dims(image, axis=0) + image = torch.from_numpy(image) + image = image.to(self.device) + return image + + def extract_id(self, np_source, f5p): + Xs = warp_and_crop_face( + np_source, + f5p, + reference_pts=get_reference_facial_points(default_square=True), + crop_size=(256, 256)) + + Xs = Image.fromarray(Xs) + Xs = self.test_transform(Xs) + Xs = Xs.unsqueeze(0).to(self.device) + with torch.no_grad(): + embeds, Xs_feats = self.arcface( + F.interpolate( + Xs, (112, 112), mode='bilinear', align_corners=True)) + return embeds, Xs + + def detect_face(self, img): + src_h, src_w, _ = img.shape + boxes, landmarks, _ = self.facer.run(img) + if boxes.shape[0] == 0: + return None + elif boxes.shape[0] > 1: + max_area = 0 + max_index = 0 + for i in range(boxes.shape[0]): + bbox_width = boxes[i][2] - boxes[i][0] + bbox_height = boxes[i][3] - boxes[i][1] + area = int(bbox_width) * int(bbox_height) + if area > max_area: + max_index = i + max_area = area + return landmarks[max_index] + else: + return landmarks[0] + + def compute_3d_params(self, Xs, Xt): + kp_fuse = {} + kp_t = {} + + c_s = self.f_3d( + F.interpolate(Xs * 0.5 + 0.5, size=224, mode='bilinear')) + c_t = self.f_3d( + F.interpolate(Xt * 0.5 + 0.5, size=224, mode='bilinear')) + c_fuse = torch.cat(((c_s[:, :80] + c_t[:, :80]) / 2, c_t[:, 80:]), + dim=1) + _, _, _, q_fuse = self.face_model.compute_for_render(c_fuse) + q_fuse = q_fuse / 224 + q_fuse[..., 1] = 1 - q_fuse[..., 1] + q_fuse = q_fuse * 2 - 1 + delta = int((17 - self.num_kp) / 2) + + _, _, _, q_t = self.face_model.compute_for_render(c_t) + q_t = q_t / 224 + q_t[..., 1] = 1 - q_t[..., 1] + q_t = q_t * 2 - 1 + + kp_fuse['value'] = q_fuse[:, delta:17 - delta, :] + kp_t['value'] = q_t[:, delta:17 - delta, :] + + return kp_fuse, kp_t + + def inference(self, template_img, user_img): + ori_h, ori_w, _ = template_img.shape + + template_img = template_img.cpu().numpy() + user_img = user_img.cpu().numpy() + + user_img_bgr = user_img[:, :, ::-1] + landmark_source = self.detect_face(user_img) + if landmark_source is None: + logger.warning('No face detected in user image!') + return template_img + f5p_user = get_f5p(landmark_source, user_img_bgr) + + template_img_bgr = template_img[:, :, ::-1] + landmark_template = self.detect_face(template_img) + if landmark_template is None: + logger.warning('No face detected in template image!') + return template_img + f5p_template = get_f5p(landmark_template, template_img_bgr) + + Xs_embeds, Xs = self.extract_id(user_img, f5p_user) + Xt, trans_inv = warp_and_crop_face( + template_img, + f5p_template, + reference_pts=get_reference_facial_points(default_square=True), + crop_size=(256, 256), + return_trans_inv=True) + + trans_inv = trans_inv.astype(np.float32) + trans_inv = torch.from_numpy(trans_inv) + trans_inv = trans_inv.to(self.device) + Xt_raw = self.image_transform(template_img, is_norm=False) + Xt = self.image_transform(Xt) + + with torch.no_grad(): + kp_fuse, kp_t = self.compute_3d_params(Xs, Xt) + Yt, _, _ = self.netG(Xt, Xs_embeds, kp_fuse, kp_t) + Yt = self.ganwrap.process_tensor(Yt) + Yt = Yt * 0.5 + 0.5 + Yt = torch.clamp(Yt, 0, 1) + + Yt_trans_inv = warp_affine_torch(Yt, trans_inv, (ori_h, ori_w)) + mask_ = warp_affine_torch(self.mask, trans_inv, (ori_h, ori_w)) + + Yt_trans_inv = mask_ * Yt_trans_inv + (1 - mask_) * Xt_raw + Yt_trans_inv = Yt_trans_inv.squeeze().permute(1, 2, + 0).cpu().numpy() + Yt_trans_inv = Yt_trans_inv.astype(np.float32) + out_img = Yt_trans_inv[:, :, ::-1] * 255. + + logger.info('model inference done') + + return out_img.astype(np.uint8) diff --git a/modelscope/models/cv/image_face_fusion/network/__init__.py b/modelscope/models/cv/image_face_fusion/network/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/image_face_fusion/network/aad_layer.py b/modelscope/models/cv/image_face_fusion/network/aad_layer.py new file mode 100644 index 00000000..adb68518 --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/network/aad_layer.py @@ -0,0 +1,99 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import torch +import torch.nn as nn + +from .ops import SpectralNorm + + +class AADLayer(nn.Module): + + def __init__(self, c_x, attr_c, c_id=256): + super(AADLayer, self).__init__() + self.attr_c = attr_c + self.c_id = c_id + self.c_x = c_x + + ks = 3 + pw = ks // 2 + + nhidden = 128 + + self.mlp_shared = nn.Sequential( + nn.ReflectionPad2d(pw), + nn.Conv2d(attr_c, nhidden, kernel_size=ks, padding=0), nn.ReLU()) + self.pad = nn.ReflectionPad2d(pw) + self.conv1 = nn.Conv2d( + nhidden, c_x, kernel_size=ks, stride=1, padding=0) + self.conv2 = nn.Conv2d( + nhidden, c_x, kernel_size=ks, stride=1, padding=0) + self.fc1 = nn.Linear(c_id, c_x) + self.fc2 = nn.Linear(c_id, c_x) + + self.norm = PositionalNorm2d + + self.pad_h = nn.ReflectionPad2d(pw) + self.conv_h = nn.Conv2d(c_x, 1, kernel_size=ks, stride=1, padding=0) + + def forward(self, h_in, z_attr, z_id): + + h = self.norm(h_in) + actv = self.mlp_shared(z_attr) + gamma_attr = self.conv1(self.pad(actv)) + beta_attr = self.conv2(self.pad(actv)) + + gamma_id = self.fc1(z_id) + beta_id = self.fc2(z_id) + A = gamma_attr * h + beta_attr + gamma_id = gamma_id.reshape(h.shape[0], self.c_x, 1, 1).expand_as(h) + beta_id = beta_id.reshape(h.shape[0], self.c_x, 1, 1).expand_as(h) + B = gamma_id * h + beta_id + + M = torch.sigmoid(self.conv_h(self.pad_h(h))) + + out = (torch.ones_like(M).to(M.device) - M) * A + M * B + + return out + + +def PositionalNorm2d(x, epsilon=1e-5): + mean = x.mean(dim=1, keepdim=True) + std = x.var(dim=1, keepdim=True).add(epsilon).sqrt() + output = (x - mean) / std + return output + + +class AAD_ResBlk(nn.Module): + + def __init__(self, cin, cout, c_attr, c_id=256): + super(AAD_ResBlk, self).__init__() + self.cin = cin + self.cout = cout + self.learned_shortcut = (self.cin != self.cout) + fmiddle = min(self.cin, self.cout) + + self.AAD1 = AADLayer(cin, c_attr, c_id) + self.AAD2 = AADLayer(fmiddle, c_attr, c_id) + self.pad = nn.ReflectionPad2d(1) + self.conv1 = SpectralNorm( + nn.Conv2d(cin, fmiddle, kernel_size=3, stride=1, padding=0)) + self.conv2 = SpectralNorm( + nn.Conv2d(fmiddle, cout, kernel_size=3, stride=1, padding=0)) + + self.relu1 = nn.LeakyReLU(2e-1) + self.relu2 = nn.LeakyReLU(2e-1) + + if self.learned_shortcut: + self.AAD3 = AADLayer(cin, c_attr, c_id) + self.conv3 = SpectralNorm( + nn.Conv2d(cin, cout, kernel_size=1, bias=False)) + + def forward(self, h, z_attr, z_id): + x = self.conv1(self.pad(self.relu1(self.AAD1(h, z_attr, z_id)))) + x = self.conv2(self.pad(self.relu2(self.AAD2(x, z_attr, z_id)))) + + if self.learned_shortcut: + h = self.conv3(self.AAD3(h, z_attr, z_id)) + + x = x + h + + return x diff --git a/modelscope/models/cv/image_face_fusion/network/aei_flow_net.py b/modelscope/models/cv/image_face_fusion/network/aei_flow_net.py new file mode 100644 index 00000000..a2047671 --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/network/aei_flow_net.py @@ -0,0 +1,251 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .aad_layer import AAD_ResBlk +from .dense_motion import DenseMotionNetwork +from .ops import SpectralNorm, init_func + + +class Conv4x4(nn.Module): + + def __init__(self, in_c, out_c): + super(Conv4x4, self).__init__() + self.conv = nn.Conv2d( + in_channels=in_c, + out_channels=out_c, + kernel_size=4, + stride=2, + padding=1, + bias=False) + self.norm = nn.BatchNorm2d(out_c) + self.lrelu = nn.LeakyReLU(0.1) + + def forward(self, feat): + x = self.conv(feat) + x = self.norm(x) + x = self.lrelu(x) + return x + + +class DeConv4x4(nn.Module): + + def __init__(self, in_c, out_c): + super(DeConv4x4, self).__init__() + self.deconv = nn.ConvTranspose2d( + in_channels=in_c, + out_channels=out_c, + kernel_size=4, + stride=2, + padding=1, + bias=False) + self.bn = nn.BatchNorm2d(out_c) + self.lrelu = nn.LeakyReLU(0.1) + + def forward(self, input, skip): + x = self.deconv(input) + x = self.bn(x) + x = self.lrelu(x) + return torch.cat((x, skip), dim=1) + + +class Attention(nn.Module): + + def __init__(self, ch, use_sn=True): + super(Attention, self).__init__() + self.ch = ch + self.theta = nn.Conv2d( + self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False) + self.phi = nn.Conv2d( + self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False) + self.g = nn.Conv2d( + self.ch, self.ch // 2, kernel_size=1, padding=0, bias=False) + self.o = nn.Conv2d( + self.ch // 2, self.ch, kernel_size=1, padding=0, bias=False) + if use_sn: + self.theta = SpectralNorm(self.theta) + self.phi = SpectralNorm(self.phi) + self.g = SpectralNorm(self.g) + self.o = SpectralNorm(self.o) + self.gamma = nn.Parameter(torch.tensor(0.), requires_grad=True) + + def forward(self, x, y=None): + theta = self.theta(x) + phi = F.max_pool2d(self.phi(x), [2, 2]) + g = F.max_pool2d(self.g(x), [2, 2]) + theta = theta.view(-1, self.ch // 8, x.shape[2] * x.shape[3]) + phi = phi.view(-1, self.ch // 8, x.shape[2] * x.shape[3] // 4) + g = g.view(-1, self.ch // 2, x.shape[2] * x.shape[3] // 4) + beta = F.softmax(torch.bmm(theta.transpose(1, 2), phi), -1) + o = self.o( + torch.bmm(g, beta.transpose(1, 2)).view(-1, self.ch // 2, + x.shape[2], x.shape[3])) + return self.gamma * o + x + + +class MLAttrEncoder(nn.Module): + + def __init__(self): + super(MLAttrEncoder, self).__init__() + self.conv1 = Conv4x4(3, 32) + self.conv2 = Conv4x4(32, 64) + self.conv3 = Conv4x4(64, 128) + self.conv4 = Conv4x4(128, 256) + self.conv5 = Conv4x4(256, 512) + self.conv6 = Conv4x4(512, 1024) + self.conv7 = Conv4x4(1024, 1024) + + self.deconv1 = DeConv4x4(1024, 1024) + self.deconv2 = DeConv4x4(2048, 512) + self.deconv3 = DeConv4x4(1024, 256) + self.deconv4 = DeConv4x4(512, 128) + self.deconv5 = DeConv4x4(256, 64) + self.deconv6 = DeConv4x4(128, 32) + + self.apply(init_func) + + def forward(self, Xt): + feat1 = self.conv1(Xt) + feat2 = self.conv2(feat1) + feat3 = self.conv3(feat2) + feat4 = self.conv4(feat3) + feat5 = self.conv5(feat4) + feat6 = self.conv6(feat5) + z_attr1 = self.conv7(feat6) + + z_attr2 = self.deconv1(z_attr1, feat6) + z_attr3 = self.deconv2(z_attr2, feat5) + z_attr4 = self.deconv3(z_attr3, feat4) + z_attr5 = self.deconv4(z_attr4, feat3) + z_attr6 = self.deconv5(z_attr5, feat2) + z_attr7 = self.deconv6(z_attr6, feat1) + z_attr8 = F.interpolate( + z_attr7, scale_factor=2, mode='bilinear', align_corners=True) + + return z_attr1, z_attr2, z_attr3, z_attr4, z_attr5, z_attr6, z_attr7, z_attr8 + + +class AADGenerator(nn.Module): + + def __init__(self, c_id=256): + super(AADGenerator, self).__init__() + self.up1 = nn.ConvTranspose2d( + c_id, 1024, kernel_size=2, stride=1, padding=0) + self.AADBlk1 = AAD_ResBlk(1024, 1024, 1024, c_id) + self.AADBlk2 = AAD_ResBlk(1024, 1024, 2048, c_id) + self.AADBlk3 = AAD_ResBlk(1024, 1024, 1024, c_id) + self.AADBlk4 = AAD_ResBlk(1024, 512, 512, c_id) + self.AADBlk5 = AAD_ResBlk(512, 256, 256, c_id) + self.AADBlk6 = AAD_ResBlk(256, 128, 128, c_id) + self.AADBlk7 = AAD_ResBlk(128, 64, 64, c_id) + self.AADBlk8 = AAD_ResBlk(64, 3, 64, c_id) + + self.sa = Attention(512, use_sn=True) + + self.apply(init_func) + + def forward(self, z_attr, z_id, deformation): + + m = self.up1(z_id.reshape(z_id.shape[0], -1, 1, 1)) + m2 = F.interpolate( + self.AADBlk1(m, z_attr[0], z_id), + scale_factor=2, + mode='bilinear', + align_corners=True) + m3 = F.interpolate( + self.AADBlk2(m2, z_attr[1], z_id), + scale_factor=2, + mode='bilinear', + align_corners=True) + m4 = F.interpolate( + self.AADBlk3(m3, z_attr[2], z_id), + scale_factor=2, + mode='bilinear', + align_corners=True) + m5 = F.interpolate( + self.AADBlk4(m4, z_attr[3], z_id), + scale_factor=2, + mode='bilinear', + align_corners=True) + m5 = self.sa(m5) + m6 = F.interpolate( + self.AADBlk5(m5, z_attr[4], z_id), + scale_factor=2, + mode='bilinear', + align_corners=True) + m7 = F.interpolate( + self.AADBlk6(m6, z_attr[5], z_id), + scale_factor=2, + mode='bilinear', + align_corners=True) + m8 = F.interpolate( + self.AADBlk7(m7, z_attr[6], z_id), + scale_factor=2, + mode='bilinear', + align_corners=True) + + y = self.AADBlk8(m8, z_attr[7], z_id) + + return torch.tanh(y) + + def deform_input(self, inp, deformation): + _, h_old, w_old, _ = deformation.shape + _, _, h, w = inp.shape + if h_old != h or w_old != w: + deformation = deformation.permute(0, 3, 1, 2) + deformation = F.interpolate( + deformation, size=(h, w), mode='bilinear') + deformation = deformation.permute(0, 2, 3, 1) + return F.grid_sample(inp, deformation), deformation + + +class AEI_Net(nn.Module): + + def __init__(self, c_id=256, num_kp=17, device=torch.device('cuda')): + super(AEI_Net, self).__init__() + self.device = device + self.encoder = MLAttrEncoder() + self.generator = AADGenerator(c_id) + self.dense_motion_network = DenseMotionNetwork( + num_kp=num_kp, num_channels=3, estimate_occlusion_map=False) + + def deform_input(self, inp, deformation): + _, h_old, w_old, _ = deformation.shape + _, _, h, w = inp.shape + if h_old != h or w_old != w: + deformation = deformation.permute(0, 3, 1, 2) + deformation = F.interpolate( + deformation, size=(h, w), mode='bilinear') + deformation = deformation.permute(0, 2, 3, 1) + return F.grid_sample(inp, deformation), deformation + + def flow_change(self, x, flow): + n, c, h, w = x.size() + yv, xv = torch.meshgrid([torch.arange(h), torch.arange(w)]) + xv = xv.float() / (w - 1) * 2.0 - 1 + yv = yv.float() / (h - 1) * 2.0 - 1 + grid = torch.cat((xv.unsqueeze(-1), yv.unsqueeze(-1)), + -1).unsqueeze(0).to(self.device) + flow_delta = flow - grid + return flow_delta + + def forward(self, Xt, z_id, kp_fuse, kp_t): + output_flow = {} + dense_motion = self.dense_motion_network( + source_image=Xt, kp_driving=kp_fuse, kp_source=kp_t) + deformation = dense_motion['deformation'] + + with torch.no_grad(): + Xt_warp, _ = self.deform_input(Xt, deformation) + attr = self.encoder(Xt_warp) + + Y = self.generator(attr, z_id, deformation) + + output_flow['deformed'], flow = self.deform_input(Xt, deformation) + output_flow['flow'] = self.flow_change(Xt, flow) + + return Y, attr, output_flow + + def get_attr(self, X): + return self.encoder(X) diff --git a/modelscope/models/cv/image_face_fusion/network/bfm.py b/modelscope/models/cv/image_face_fusion/network/bfm.py new file mode 100644 index 00000000..fe6a42e5 --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/network/bfm.py @@ -0,0 +1,249 @@ +# The implementation is adopted from Deep3DFaceRecon_pytorch, made publicly available under the MIT License +# at https://github.com/sicxu/Deep3DFaceRecon_pytorch/blob/master/models/bfm.py +import os + +import numpy as np +import torch +import torch.nn.functional as F +from scipy.io import loadmat + + +def perspective_projection(focal, center): + return np.array([focal, 0, center, 0, focal, center, 0, 0, + 1]).reshape([3, 3]).astype(np.float32).transpose() + + +class SH: + + def __init__(self): + self.a = [np.pi, 2 * np.pi / np.sqrt(3.), 2 * np.pi / np.sqrt(8.)] + self.c = [ + 1 / np.sqrt(4 * np.pi), + np.sqrt(3.) / np.sqrt(4 * np.pi), + 3 * np.sqrt(5.) / np.sqrt(12 * np.pi) + ] + + +class ParametricFaceModel(): + + def __init__(self, + bfm_folder='./BFM', + recenter=True, + camera_distance=10., + init_lit=np.array([0.8, 0, 0, 0, 0, 0, 0, 0, 0]), + focal=1015., + center=112., + is_train=True, + default_name='BFM_model_front.mat'): + + model = loadmat(os.path.join(bfm_folder, default_name)) + # mean face shape. [3*N,1] + self.mean_shape = model['meanshape'].astype(np.float32) + # identity basis. [3*N,80] + self.id_base = model['idBase'].astype(np.float32) + # expression basis. [3*N,64] + self.exp_base = model['exBase'].astype(np.float32) + # mean face texture. [3*N,1] (0-255) + self.mean_tex = model['meantex'].astype(np.float32) + # texture basis. [3*N,80] + self.tex_base = model['texBase'].astype(np.float32) + # face indices for each vertex that lies in. starts from 0. [N,8] + self.point_buf = model['point_buf'].astype(np.int64) - 1 + # vertex indices for each face. starts from 0. [F,3] + self.face_buf = model['tri'].astype(np.int64) - 1 + # vertex indices for 68 landmarks. starts from 0. [68,1] + self.keypoints = np.squeeze(model['keypoints']).astype(np.int64) - 1 + + if is_train: + # vertex indices for small face region to compute photometric error. starts from 0. + self.front_mask = np.squeeze(model['frontmask2_idx']).astype( + np.int64) - 1 + # vertex indices for each face from small face region. starts from 0. [f,3] + self.front_face_buf = model['tri_mask2'].astype(np.int64) - 1 + # vertex indices for pre-defined skin region to compute reflectance loss + self.skin_mask = np.squeeze(model['skinmask']) + + if recenter: + mean_shape = self.mean_shape.reshape([-1, 3]) + mean_shape = mean_shape - np.mean( + mean_shape, axis=0, keepdims=True) + self.mean_shape = mean_shape.reshape([-1, 1]) + + self.persc_proj = perspective_projection(focal, center) + self.device = 'cpu' + self.camera_distance = camera_distance + self.SH = SH() + self.init_lit = init_lit.reshape([1, 1, -1]).astype(np.float32) + + def to(self, device): + self.device = device + for key, value in self.__dict__.items(): + if type(value).__module__ == np.__name__: + setattr(self, key, torch.tensor(value).to(device)) + + def compute_shape(self, id_coeff, exp_coeff): + """ + Return: + face_shape -- torch.tensor, size (B, N, 3) + + Parameters: + id_coeff -- torch.tensor, size (B, 80), identity coeffs + exp_coeff -- torch.tensor, size (B, 64), expression coeffs + """ + batch_size = id_coeff.shape[0] + id_part = torch.einsum('ij,aj->ai', self.id_base, id_coeff) + exp_part = torch.einsum('ij,aj->ai', self.exp_base, exp_coeff) + face_shape = id_part + exp_part + self.mean_shape.reshape([1, -1]) + return face_shape.reshape([batch_size, -1, 3]) + + def compute_texture(self, tex_coeff, normalize=True): + """ + Return: + face_texture -- torch.tensor, size (B, N, 3), in RGB order, range (0, 1.) + + Parameters: + tex_coeff -- torch.tensor, size (B, 80) + """ + batch_size = tex_coeff.shape[0] + face_texture = torch.einsum('ij,aj->ai', self.tex_base, + tex_coeff) + self.mean_tex + if normalize: + face_texture = face_texture / 255. + return face_texture.reshape([batch_size, -1, 3]) + + def compute_norm(self, face_shape): + """ + Return: + vertex_norm -- torch.tensor, size (B, N, 3) + + Parameters: + face_shape -- torch.tensor, size (B, N, 3) + """ + + v1 = face_shape[:, self.face_buf[:, 0]] + v2 = face_shape[:, self.face_buf[:, 1]] + v3 = face_shape[:, self.face_buf[:, 2]] + e1 = v1 - v2 + e2 = v2 - v3 + face_norm = torch.cross(e1, e2, dim=-1) + face_norm = F.normalize(face_norm, dim=-1, p=2) + face_norm = torch.cat( + [face_norm, + torch.zeros(face_norm.shape[0], 1, 3).to(self.device)], + dim=1) + + vertex_norm = torch.sum(face_norm[:, self.point_buf], dim=2) + vertex_norm = F.normalize(vertex_norm, dim=-1, p=2) + return vertex_norm + + def compute_color(self, face_texture, face_norm, gamma): + batch_size = gamma.shape[0] + a, c = self.SH.a, self.SH.c + gamma = gamma.reshape([batch_size, 3, 9]) + gamma = gamma + self.init_lit + gamma = gamma.permute(0, 2, 1) + face_norm_p1 = face_norm[..., :1] + face_norm_p2 = face_norm[..., 1:2] + face_norm_p3 = face_norm[..., 2:] + face_norm_diff = face_norm_p1**2 - face_norm_p2**2 + temp = [ + a[0] * c[0] * torch.ones_like(face_norm_p1).to(self.device), + -a[1] * c[1] * face_norm_p2, a[1] * c[1] * face_norm_p3, + -a[1] * c[1] * face_norm_p1, + a[2] * c[2] * face_norm_p1 * face_norm_p2, + -a[2] * c[2] * face_norm_p2 * face_norm_p3, + 0.5 * a[2] * c[2] / np.sqrt(3.) * (3 * face_norm_p3**2 - 1), + -a[2] * c[2] * face_norm_p1 * face_norm_p3, + 0.5 * a[2] * c[2] * face_norm_diff + ] + + Y = torch.cat(temp, dim=-1) + r = Y @ gamma[..., :1] + g = Y @ gamma[..., 1:2] + b = Y @ gamma[..., 2:] + face_color = torch.cat([r, g, b], dim=-1) * face_texture + return face_color + + def compute_rotation(self, angles): + batch_size = angles.shape[0] + ones = torch.ones([batch_size, 1]).to(self.device) + zeros = torch.zeros([batch_size, 1]).to(self.device) + x, y, z = angles[:, :1], angles[:, 1:2], angles[:, 2:], + + temp_x = [ + ones, zeros, zeros, zeros, + torch.cos(x), -torch.sin(x), zeros, + torch.sin(x), + torch.cos(x) + ] + rot_x = torch.cat(temp_x, dim=1).reshape([batch_size, 3, 3]) + + temp_y = [ + torch.cos(y), zeros, + torch.sin(y), zeros, ones, zeros, -torch.sin(y), zeros, + torch.cos(y) + ] + rot_y = torch.cat(temp_y, dim=1).reshape([batch_size, 3, 3]) + + temp_z = [ + torch.cos(z), -torch.sin(z), zeros, + torch.sin(z), + torch.cos(z), zeros, zeros, zeros, ones + ] + rot_z = torch.cat(temp_z, dim=1).reshape([batch_size, 3, 3]) + + rot = rot_z @ rot_y @ rot_x + return rot.permute(0, 2, 1) + + def to_camera(self, face_shape): + face_shape[..., -1] = self.camera_distance - face_shape[..., -1] + return face_shape + + def to_image(self, face_shape): + # to image_plane + face_proj = face_shape @ self.persc_proj + face_proj = face_proj[..., :2] / face_proj[..., 2:] + + return face_proj + + def transform(self, face_shape, rot, trans): + return face_shape @ rot + trans.unsqueeze(1) + + def get_landmarks(self, face_proj): + return face_proj[:, self.keypoints] + + def split_coeff(self, coeffs): + id_coeffs = coeffs[:, :80] + exp_coeffs = coeffs[:, 80:144] + tex_coeffs = coeffs[:, 144:224] + angles = coeffs[:, 224:227] + gammas = coeffs[:, 227:254] + translations = coeffs[:, 254:] + return { + 'id': id_coeffs, + 'exp': exp_coeffs, + 'tex': tex_coeffs, + 'angle': angles, + 'gamma': gammas, + 'trans': translations + } + + def compute_for_render(self, coeffs): + coef_dict = self.split_coeff(coeffs) + face_shape = self.compute_shape(coef_dict['id'], coef_dict['exp']) + rotation = self.compute_rotation(coef_dict['angle']) + + face_shape_transformed = self.transform(face_shape, rotation, + coef_dict['trans']) + face_vertex = self.to_camera(face_shape_transformed) + + face_proj = self.to_image(face_vertex) + landmark = self.get_landmarks(face_proj) + + face_texture = self.compute_texture(coef_dict['tex']) + face_norm = self.compute_norm(face_shape) + face_norm_roted = face_norm @ rotation + face_color = self.compute_color(face_texture, face_norm_roted, + coef_dict['gamma']) + + return face_vertex, face_texture, face_color, landmark diff --git a/modelscope/models/cv/image_face_fusion/network/dense_motion.py b/modelscope/models/cv/image_face_fusion/network/dense_motion.py new file mode 100644 index 00000000..ba5cc2db --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/network/dense_motion.py @@ -0,0 +1,376 @@ +# The implementation is adopted from first-order-model, made publicly available under the MIT License +# at https://github.com/AliaksandrSiarohin/first-order-model/blob/master/modules/dense_motion.py + +import torch +import torch.nn.functional as F +from torch import nn + + +def kp2gaussian(kp, spatial_size, kp_variance): + """ + Transform a keypoint into gaussian like representation + """ + mean = kp['value'] + + coordinate_grid = make_coordinate_grid(spatial_size, mean.type()) + number_of_leading_dimensions = len(mean.shape) - 1 + shape = (1, ) * number_of_leading_dimensions + coordinate_grid.shape + coordinate_grid = coordinate_grid.view(*shape) + repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1) + coordinate_grid = coordinate_grid.repeat(*repeats) + + # Preprocess kp shape + shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 2) + mean = mean.view(*shape) + + mean_sub = (coordinate_grid - mean) + + out = torch.exp(-0.5 * (mean_sub**2).sum(-1) / kp_variance) + + return out + + +def make_coordinate_grid(spatial_size, type): + """ + Create a meshgrid [-1,1] x [-1,1] of given spatial_size. + """ + h, w = spatial_size + x = torch.arange(w).type(type) + y = torch.arange(h).type(type) + + x = (2 * (x / (w - 1)) - 1) + y = (2 * (y / (h - 1)) - 1) + + yy = y.view(-1, 1).repeat(1, w) + xx = x.view(1, -1).repeat(h, 1) + + meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2) + + return meshed + + +class UpBlock2d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, + in_features, + out_features, + kernel_size=3, + padding=1, + groups=1): + super(UpBlock2d, self).__init__() + + self.conv = nn.Conv2d( + in_channels=in_features, + out_channels=out_features, + kernel_size=kernel_size, + padding=padding, + groups=groups) + self.norm = nn.BatchNorm2d(out_features, affine=True) + + def forward(self, x): + out = F.interpolate(x, scale_factor=2) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class DownBlock2d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, + in_features, + out_features, + kernel_size=3, + padding=1, + groups=1): + super(DownBlock2d, self).__init__() + self.conv = nn.Conv2d( + in_channels=in_features, + out_channels=out_features, + kernel_size=kernel_size, + padding=padding, + groups=groups) + self.norm = nn.BatchNorm2d(out_features, affine=True) + self.pool = nn.AvgPool2d(kernel_size=(2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class Encoder(nn.Module): + """ + Hourglass Encoder + """ + + def __init__(self, + block_expansion, + in_features, + num_blocks=3, + max_features=256): + super(Encoder, self).__init__() + + down_blocks = [] + for i in range(num_blocks): + down_blocks.append( + DownBlock2d( + in_features if i == 0 else min(max_features, + block_expansion * (2**i)), + min(max_features, block_expansion * (2**(i + 1))), + kernel_size=3, + padding=1)) + self.down_blocks = nn.ModuleList(down_blocks) + + def forward(self, x): + outs = [x] + for down_block in self.down_blocks: + outs.append(down_block(outs[-1])) + return outs + + +class Decoder(nn.Module): + """ + Hourglass Decoder + """ + + def __init__(self, + block_expansion, + in_features, + num_blocks=3, + max_features=256): + super(Decoder, self).__init__() + + up_blocks = [] + + for i in range(num_blocks)[::-1]: + in_filters = (1 if i == num_blocks - 1 else 2) * min( + max_features, block_expansion * (2**(i + 1))) + out_filters = min(max_features, block_expansion * (2**i)) + up_blocks.append( + UpBlock2d(in_filters, out_filters, kernel_size=3, padding=1)) + + self.up_blocks = nn.ModuleList(up_blocks) + self.out_filters = block_expansion + in_features + + def forward(self, x): + out = x.pop() + for up_block in self.up_blocks: + out = up_block(out) + skip = x.pop() + out = torch.cat([out, skip], dim=1) + return out + + +class Hourglass(nn.Module): + """ + Hourglass architecture. + """ + + def __init__(self, + block_expansion, + in_features, + num_blocks=3, + max_features=256): + super(Hourglass, self).__init__() + self.encoder = Encoder(block_expansion, in_features, num_blocks, + max_features) + self.decoder = Decoder(block_expansion, in_features, num_blocks, + max_features) + self.out_filters = self.decoder.out_filters + + def forward(self, x): + return self.decoder(self.encoder(x)) + + +class AntiAliasInterpolation2d(nn.Module): + """ + Band-limited downsampling, for better preservation of the input signal. + """ + + def __init__(self, channels, scale): + super(AntiAliasInterpolation2d, self).__init__() + sigma = (1 / scale - 1) / 2 + kernel_size = 2 * round(sigma * 4) + 1 + self.ka = kernel_size // 2 + self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka + + kernel_size = [kernel_size, kernel_size] + sigma = [sigma, sigma] + # The gaussian kernel is the product of the + # gaussian function of each dimension. + kernel = 1 + meshgrids = torch.meshgrid( + [torch.arange(size, dtype=torch.float32) for size in kernel_size]) + for size, std, mgrid in zip(kernel_size, sigma, meshgrids): + mean = (size - 1) / 2 + kernel *= torch.exp(-(mgrid - mean)**2 / (2 * std**2)) + + # Make sure sum of values in gaussian kernel equals 1. + kernel = kernel / torch.sum(kernel) + # Reshape to depthwise convolutional weight + kernel = kernel.view(1, 1, *kernel.size()) + kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1)) + + self.register_buffer('weight', kernel) + self.groups = channels + self.scale = scale + inv_scale = 1 / scale + self.int_inv_scale = int(inv_scale) + + def forward(self, input): + if self.scale == 1.0: + return input + + out = F.pad(input, (self.ka, self.kb, self.ka, self.kb)) + out = F.conv2d(out, weight=self.weight, groups=self.groups) + out = out[:, :, ::self.int_inv_scale, ::self.int_inv_scale] + + return out + + +class DenseMotionNetwork(nn.Module): + """ + Module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving + """ + + def __init__(self, + num_kp, + num_channels, + estimate_occlusion_map=False, + kp_variance=0.01): + super(DenseMotionNetwork, self).__init__() + block_expansion = 64 + num_blocks = 5 + max_features = 1024 + scale_factor = 0.25 + + self.hourglass = Hourglass( + block_expansion=block_expansion, + in_features=(num_kp + 1) * (num_channels + 1), + max_features=max_features, + num_blocks=num_blocks) + + self.mask = nn.Conv2d( + self.hourglass.out_filters, + num_kp + 1, + kernel_size=(7, 7), + padding=(3, 3)) + + if estimate_occlusion_map: + self.occlusion = nn.Conv2d( + self.hourglass.out_filters, + 1, + kernel_size=(7, 7), + padding=(3, 3)) + else: + self.occlusion = None + + self.num_kp = num_kp + self.scale_factor = scale_factor + self.kp_variance = kp_variance + + if self.scale_factor != 1: + self.down = AntiAliasInterpolation2d(num_channels, + self.scale_factor) + + def create_heatmap_representations(self, source_image, kp_driving, + kp_source): + """ + Eq 6. in the paper H_k(z) + """ + spatial_size = source_image.shape[2:] + gaussian_driving = kp2gaussian( + kp_driving, + spatial_size=spatial_size, + kp_variance=self.kp_variance) + gaussian_source = kp2gaussian( + kp_source, spatial_size=spatial_size, kp_variance=self.kp_variance) + heatmap = gaussian_driving - gaussian_source + + zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], + spatial_size[1]).type(heatmap.type()) + heatmap = torch.cat([zeros, heatmap], dim=1) + heatmap = heatmap.unsqueeze(2) + return heatmap + + def create_sparse_motions(self, source_image, kp_driving, kp_source): + """ + Eq 4. in the paper T_{s<-d}(z) + """ + bs, _, h, w = source_image.shape + identity_grid = make_coordinate_grid((h, w), + type=kp_source['value'].type()) + identity_grid = identity_grid.view(1, 1, h, w, 2) + coordinate_grid = identity_grid - kp_driving['value'].view( + bs, self.num_kp, 1, 1, 2) + if 'jacobian' in kp_driving: + jacobian = torch.matmul(kp_source['jacobian'], + torch.inverse(kp_driving['jacobian'])) + jacobian = jacobian.unsqueeze(-3).unsqueeze(-3) + jacobian = jacobian.repeat(1, 1, h, w, 1, 1) + coordinate_grid = torch.matmul(jacobian, + coordinate_grid.unsqueeze(-1)) + coordinate_grid = coordinate_grid.squeeze(-1) + + driving_to_source = coordinate_grid + kp_source['value'].view( + bs, self.num_kp, 1, 1, 2) + + identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1) + sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1) + return sparse_motions + + def create_deformed_source_image(self, source_image, sparse_motions): + bs, _, h, w = source_image.shape + source_repeat = source_image.unsqueeze(1).unsqueeze(1).repeat( + 1, self.num_kp + 1, 1, 1, 1, 1) + temp_dim = bs * (self.num_kp + 1) + source_repeat = source_repeat.view(temp_dim, -1, h, w) + sparse_motions = sparse_motions.view((temp_dim, h, w, -1)) + sparse_deformed = F.grid_sample(source_repeat, sparse_motions) + sparse_deformed = sparse_deformed.view((bs, self.num_kp + 1, -1, h, w)) + return sparse_deformed + + def forward(self, source_image, kp_driving, kp_source): + if self.scale_factor != 1: + source_image = self.down(source_image) + + bs, _, h, w = source_image.shape + + out_dict = dict() + heatmap_representation = self.create_heatmap_representations( + source_image, kp_driving, kp_source) + sparse_motion = self.create_sparse_motions(source_image, kp_driving, + kp_source) + deformed_source = self.create_deformed_source_image( + source_image, sparse_motion) + out_dict['sparse_deformed'] = deformed_source + + input = torch.cat([heatmap_representation, deformed_source], dim=2) + input = input.view(bs, -1, h, w) + + prediction = self.hourglass(input) + + mask = self.mask(prediction) + mask = F.softmax(mask, dim=1) + out_dict['mask'] = mask + mask = mask.unsqueeze(2) + sparse_motion = sparse_motion.permute(0, 1, 4, 2, 3) + deformation = (sparse_motion * mask).sum(dim=1) + deformation = deformation.permute(0, 2, 3, 1) + + out_dict['deformation'] = deformation + + if self.occlusion: + occlusion_map = torch.sigmoid(self.occlusion(prediction)) + out_dict['occlusion_map'] = occlusion_map + + return out_dict diff --git a/modelscope/models/cv/image_face_fusion/network/facerecon_model.py b/modelscope/models/cv/image_face_fusion/network/facerecon_model.py new file mode 100644 index 00000000..53b34d3c --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/network/facerecon_model.py @@ -0,0 +1,559 @@ +# The implementation is adopted from Deep3DFaceRecon_pytorch, made publicly available under the MIT License +# at https://github.com/sicxu/Deep3DFaceRecon_pytorch/blob/master/models/networks.py + +import os +from typing import Any, Callable, List, Optional, Type, Union + +import torch +import torch.nn as nn +from torch import Tensor +from torch.optim import lr_scheduler + +try: + from torch.hub import load_state_dict_from_url +except ImportError: + from torch.utils.model_zoo import load_url as load_state_dict_from_url + + +def filter_state_dict(state_dict, remove_name='fc'): + new_state_dict = {} + for key in state_dict: + if remove_name in key: + continue + new_state_dict[key] = state_dict[key] + return new_state_dict + + +def get_scheduler(optimizer, opt): + """Return a learning rate scheduler + + Parameters: + optimizer -- the optimizer of the network + opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．　 + opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine + + For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers. + See https://pytorch.org/docs/stable/optim.html for more details. + """ + if opt.lr_policy == 'linear': + + def lambda_rule(epoch): + lr_l = 1.0 - max(0, epoch + opt.epoch_count + - opt.n_epochs) / float(opt.n_epochs + 1) + return lr_l + + scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) + elif opt.lr_policy == 'step': + scheduler = lr_scheduler.StepLR( + optimizer, step_size=opt.lr_decay_epochs, gamma=0.2) + elif opt.lr_policy == 'plateau': + scheduler = lr_scheduler.ReduceLROnPlateau( + optimizer, mode='min', factor=0.2, threshold=0.01, patience=5) + elif opt.lr_policy == 'cosine': + scheduler = lr_scheduler.CosineAnnealingLR( + optimizer, T_max=opt.n_epochs, eta_min=0) + else: + return NotImplementedError( + 'learning rate policy [%s] is not implemented', opt.lr_policy) + return scheduler + + +def define_net_recon(net_recon, use_last_fc=False, init_path=None): + return ReconNetWrapper( + net_recon, use_last_fc=use_last_fc, init_path=init_path) + + +class ReconNetWrapper(nn.Module): + fc_dim = 257 + + def __init__(self, net_recon, use_last_fc=False, init_path=None): + super(ReconNetWrapper, self).__init__() + self.use_last_fc = use_last_fc + if net_recon not in func_dict: + return NotImplementedError('network [%s] is not implemented', + net_recon) + func, last_dim = func_dict[net_recon] + backbone = func(use_last_fc=use_last_fc, num_classes=self.fc_dim) + if init_path and os.path.isfile(init_path): + state_dict = filter_state_dict( + torch.load(init_path, map_location='cpu')) + backbone.load_state_dict(state_dict) + print('loading init net_recon %s from %s' % (net_recon, init_path)) + self.backbone = backbone + if not use_last_fc: + self.final_layers = nn.ModuleList([ + conv1x1(last_dim, 80, bias=True), # id layer + conv1x1(last_dim, 64, bias=True), # exp layer + conv1x1(last_dim, 80, bias=True), # tex layer + conv1x1(last_dim, 3, bias=True), # angle layer + conv1x1(last_dim, 27, bias=True), # gamma layer + conv1x1(last_dim, 2, bias=True), # tx, ty + conv1x1(last_dim, 1, bias=True) # tz + ]) + for m in self.final_layers: + nn.init.constant_(m.weight, 0.) + nn.init.constant_(m.bias, 0.) + + def forward(self, x): + x = self.backbone(x) + if not self.use_last_fc: + output = [] + for layer in self.final_layers: + output.append(layer(x)) + x = torch.flatten(torch.cat(output, dim=1), 1) + return x + + +# adapted from https://github.com/pytorch/vision/edit/master/torchvision/models/resnet.py +__all__ = [ + 'ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', + 'resnext50_32x4d', 'resnext101_32x8d', 'wide_resnet50_2', + 'wide_resnet101_2' +] + +model_urls = { + 'resnet18': + 'https://download.pytorch.org/models/resnet18-f37072fd.pth', + 'resnet34': + 'https://download.pytorch.org/models/resnet34-b627a593.pth', + 'resnet50': + 'https://download.pytorch.org/models/resnet50-0676ba61.pth', + 'resnet101': + 'https://download.pytorch.org/models/resnet101-63fe2227.pth', + 'resnet152': + 'https://download.pytorch.org/models/resnet152-394f9c45.pth', + 'resnext50_32x4d': + 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth', + 'resnext101_32x8d': + 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth', + 'wide_resnet50_2': + 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth', + 'wide_resnet101_2': + 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth', +} + + +def conv3x3(in_planes: int, + out_planes: int, + stride: int = 1, + groups: int = 1, + dilation: int = 1) -> nn.Conv2d: + """3x3 convolution with padding""" + return nn.Conv2d( + in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=dilation, + groups=groups, + bias=False, + dilation=dilation) + + +def conv1x1(in_planes: int, + out_planes: int, + stride: int = 1, + bias: bool = False) -> nn.Conv2d: + """1x1 convolution""" + return nn.Conv2d( + in_planes, out_planes, kernel_size=1, stride=stride, bias=bias) + + +class BasicBlock(nn.Module): + expansion: int = 1 + + def __init__( + self, + inplanes: int, + planes: int, + stride: int = 1, + downsample: Optional[nn.Module] = None, + groups: int = 1, + base_width: int = 64, + dilation: int = 1, + norm_layer: Optional[Callable[..., nn.Module]] = None) -> None: + super(BasicBlock, self).__init__() + if norm_layer is None: + norm_layer = nn.BatchNorm2d + if groups != 1 or base_width != 64: + raise ValueError( + 'BasicBlock only supports groups=1 and base_width=64') + if dilation > 1: + raise NotImplementedError( + 'Dilation > 1 not supported in BasicBlock') + # Both self.conv1 and self.downsample layers downsample the input when stride != 1 + self.conv1 = conv3x3(inplanes, planes, stride) + self.bn1 = norm_layer(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3(planes, planes) + self.bn2 = norm_layer(planes) + self.downsample = downsample + self.stride = stride + + def forward(self, x: Tensor) -> Tensor: + identity = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2) + # while original implementation places the stride at the first 1x1 convolution(self.conv1) + # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385. + # This variant is also known as ResNet V1.5 and improves accuracy according to + # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch. + + expansion: int = 4 + + def __init__( + self, + inplanes: int, + planes: int, + stride: int = 1, + downsample: Optional[nn.Module] = None, + groups: int = 1, + base_width: int = 64, + dilation: int = 1, + norm_layer: Optional[Callable[..., nn.Module]] = None) -> None: + super(Bottleneck, self).__init__() + if norm_layer is None: + norm_layer = nn.BatchNorm2d + width = int(planes * (base_width / 64.)) * groups + # Both self.conv2 and self.downsample layers downsample the input when stride != 1 + self.conv1 = conv1x1(inplanes, width) + self.bn1 = norm_layer(width) + self.conv2 = conv3x3(width, width, stride, groups, dilation) + self.bn2 = norm_layer(width) + self.conv3 = conv1x1(width, planes * self.expansion) + self.bn3 = norm_layer(planes * self.expansion) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x: Tensor) -> Tensor: + identity = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu(out) + + return out + + +class ResNet(nn.Module): + + def __init__( + self, + block: Type[Union[BasicBlock, Bottleneck]], + layers: List[int], + num_classes: int = 1000, + zero_init_residual: bool = False, + use_last_fc: bool = False, + groups: int = 1, + width_per_group: int = 64, + replace_stride_with_dilation: Optional[List[bool]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None) -> None: + super(ResNet, self).__init__() + if norm_layer is None: + norm_layer = nn.BatchNorm2d + self._norm_layer = norm_layer + + self.inplanes = 64 + self.dilation = 1 + if replace_stride_with_dilation is None: + # each element in the tuple indicates if we should replace + # the 2x2 stride with a dilated convolution instead + replace_stride_with_dilation = [False, False, False] + if len(replace_stride_with_dilation) != 3: + raise ValueError('replace_stride_with_dilation should be None ' + 'or a 3-element tuple, got {}'.format( + replace_stride_with_dilation)) + self.use_last_fc = use_last_fc + self.groups = groups + self.base_width = width_per_group + self.conv1 = nn.Conv2d( + 3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False) + self.bn1 = norm_layer(self.inplanes) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer( + block, + 128, + layers[1], + stride=2, + dilate=replace_stride_with_dilation[0]) + self.layer3 = self._make_layer( + block, + 256, + layers[2], + stride=2, + dilate=replace_stride_with_dilation[1]) + self.layer4 = self._make_layer( + block, + 512, + layers[3], + stride=2, + dilate=replace_stride_with_dilation[2]) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + + if self.use_last_fc: + self.fc = nn.Linear(512 * block.expansion, num_classes) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_( + m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + # Zero-initialize the last BN in each residual branch, + # so that the residual branch starts with zeros, and each residual block behaves like an identity. + # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 + if zero_init_residual: + for m in self.modules(): + if isinstance(m, Bottleneck): + nn.init.constant_(m.bn3.weight, + 0) # type: ignore[arg-type] + elif isinstance(m, BasicBlock): + nn.init.constant_(m.bn2.weight, + 0) # type: ignore[arg-type] + + def _make_layer(self, + block: Type[Union[BasicBlock, Bottleneck]], + planes: int, + blocks: int, + stride: int = 1, + dilate: bool = False) -> nn.Sequential: + norm_layer = self._norm_layer + downsample = None + previous_dilation = self.dilation + if dilate: + self.dilation *= stride + stride = 1 + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + conv1x1(self.inplanes, planes * block.expansion, stride), + norm_layer(planes * block.expansion), + ) + + layers = [] + layers.append( + block(self.inplanes, planes, stride, downsample, self.groups, + self.base_width, previous_dilation, norm_layer)) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append( + block( + self.inplanes, + planes, + groups=self.groups, + base_width=self.base_width, + dilation=self.dilation, + norm_layer=norm_layer)) + + return nn.Sequential(*layers) + + def _forward_impl(self, x: Tensor) -> Tensor: + # See note [TorchScript super()] + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + x = self.avgpool(x) + if self.use_last_fc: + x = torch.flatten(x, 1) + x = self.fc(x) + return x + + def forward(self, x: Tensor) -> Tensor: + return self._forward_impl(x) + + +def _resnet(arch: str, block: Type[Union[BasicBlock, + Bottleneck]], layers: List[int], + pretrained: bool, progress: bool, **kwargs: Any) -> ResNet: + model = ResNet(block, layers, **kwargs) + if pretrained: + state_dict = load_state_dict_from_url( + model_urls[arch], progress=progress) + model.load_state_dict(state_dict) + return model + + +def resnet18(pretrained: bool = False, + progress: bool = True, + **kwargs: Any) -> ResNet: + r"""ResNet-18 model from + `"Deep Residual Learning for Image Recognition" `_. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, + **kwargs) + + +def resnet34(pretrained: bool = False, + progress: bool = True, + **kwargs: Any) -> ResNet: + r"""ResNet-34 model from + `"Deep Residual Learning for Image Recognition" `_. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, + **kwargs) + + +def resnet50(pretrained: bool = False, + progress: bool = True, + **kwargs: Any) -> ResNet: + r"""ResNet-50 model from + `"Deep Residual Learning for Image Recognition" `_. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, + **kwargs) + + +def resnet101(pretrained: bool = False, + progress: bool = True, + **kwargs: Any) -> ResNet: + r"""ResNet-101 model from + `"Deep Residual Learning for Image Recognition" `_. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, + progress, **kwargs) + + +def resnet152(pretrained: bool = False, + progress: bool = True, + **kwargs: Any) -> ResNet: + r"""ResNet-152 model from + `"Deep Residual Learning for Image Recognition" `_. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, + progress, **kwargs) + + +def resnext50_32x4d(pretrained: bool = False, + progress: bool = True, + **kwargs: Any) -> ResNet: + r"""ResNeXt-50 32x4d model from + `"Aggregated Residual Transformation for Deep Neural Networks" `_. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + kwargs['groups'] = 32 + kwargs['width_per_group'] = 4 + return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3], pretrained, + progress, **kwargs) + + +def resnext101_32x8d(pretrained: bool = False, + progress: bool = True, + **kwargs: Any) -> ResNet: + r"""ResNeXt-101 32x8d model from + `"Aggregated Residual Transformation for Deep Neural Networks" `_. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + kwargs['groups'] = 32 + kwargs['width_per_group'] = 8 + return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], pretrained, + progress, **kwargs) + + +def wide_resnet50_2(pretrained: bool = False, + progress: bool = True, + **kwargs: Any) -> ResNet: + r"""Wide ResNet-50-2 model from + `"Wide Residual Networks" `_. + + The model is the same as ResNet except for the bottleneck number of channels + which is twice larger in every block. The number of channels in outer 1x1 + convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 + channels, and in Wide ResNet-50-2 has 2048-1024-2048. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + kwargs['width_per_group'] = 64 * 2 + return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3], pretrained, + progress, **kwargs) + + +def wide_resnet101_2(pretrained: bool = False, + progress: bool = True, + **kwargs: Any) -> ResNet: + r"""Wide ResNet-101-2 model from + `"Wide Residual Networks" `_. + + The model is the same as ResNet except for the bottleneck number of channels + which is twice larger in every block. The number of channels in outer 1x1 + convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 + channels, and in Wide ResNet-50-2 has 2048-1024-2048. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + progress (bool): If True, displays a progress bar of the download to stderr + """ + kwargs['width_per_group'] = 64 * 2 + return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3], pretrained, + progress, **kwargs) + + +func_dict = {'resnet18': (resnet18, 512), 'resnet50': (resnet50, 2048)} diff --git a/modelscope/models/cv/image_face_fusion/network/model_irse.py b/modelscope/models/cv/image_face_fusion/network/model_irse.py new file mode 100644 index 00000000..4e212bd0 --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/network/model_irse.py @@ -0,0 +1,249 @@ +# The implementation is adopted from face.evoLVe, made publicly available under the MIT License +# at https://github.com/ZhaoJ9014/face.evoLVe/blob/master/backbone/model_irse.py +from collections import namedtuple + +import torch +import torch.nn as nn +from torch.nn import (AdaptiveAvgPool2d, BatchNorm1d, BatchNorm2d, Conv2d, + Dropout, Linear, MaxPool2d, Module, PReLU, ReLU, + Sequential, Sigmoid) + + +class Flatten(Module): + + def forward(self, input): + return input.view(input.size(0), -1) + + +def l2_norm(input, axis=1): + norm = torch.norm(input, 2, axis, True) + output = torch.div(input, norm) + + return output + + +class SEModule(Module): + + def __init__(self, channels, reduction): + super(SEModule, self).__init__() + self.avg_pool = AdaptiveAvgPool2d(1) + self.fc1 = Conv2d( + channels, + channels // reduction, + kernel_size=1, + padding=0, + bias=False) + + nn.init.xavier_uniform_(self.fc1.weight.data) + + self.relu = ReLU(inplace=True) + self.fc2 = Conv2d( + channels // reduction, + channels, + kernel_size=1, + padding=0, + bias=False) + + self.sigmoid = Sigmoid() + + def forward(self, x): + module_input = x + x = self.avg_pool(x) + x = self.fc1(x) + x = self.relu(x) + x = self.fc2(x) + x = self.sigmoid(x) + + return module_input * x + + +class bottleneck_IR(Module): + + def __init__(self, in_channel, depth, stride): + super(bottleneck_IR, self).__init__() + if in_channel == depth: + self.shortcut_layer = MaxPool2d(1, stride) + else: + self.shortcut_layer = Sequential( + Conv2d(in_channel, depth, (1, 1), stride, bias=False), + BatchNorm2d(depth)) + self.res_layer = Sequential( + BatchNorm2d(in_channel), + Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), + PReLU(depth), Conv2d(depth, depth, (3, 3), stride, 1, bias=False), + BatchNorm2d(depth)) + + def forward(self, x): + shortcut = self.shortcut_layer(x) + res = self.res_layer(x) + + return res + shortcut + + +class bottleneck_IR_SE(Module): + + def __init__(self, in_channel, depth, stride): + super(bottleneck_IR_SE, self).__init__() + if in_channel == depth: + self.shortcut_layer = MaxPool2d(1, stride) + else: + self.shortcut_layer = Sequential( + Conv2d(in_channel, depth, (1, 1), stride, bias=False), + BatchNorm2d(depth)) + self.res_layer = Sequential( + BatchNorm2d(in_channel), + Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), + PReLU(depth), Conv2d(depth, depth, (3, 3), stride, 1, bias=False), + BatchNorm2d(depth), SEModule(depth, 16)) + + def forward(self, x): + shortcut = self.shortcut_layer(x) + res = self.res_layer(x) + + return res + shortcut + + +class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])): + '''A named tuple describing a ResNet block.''' + + +def get_block(in_channel, depth, num_units, stride=2): + + return [Bottleneck(in_channel, depth, stride) + ] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)] + + +def get_blocks(num_layers): + if num_layers == 50: + blocks = [ + get_block(in_channel=64, depth=64, num_units=3), + get_block(in_channel=64, depth=128, num_units=4), + get_block(in_channel=128, depth=256, num_units=14), + get_block(in_channel=256, depth=512, num_units=3) + ] + elif num_layers == 100: + blocks = [ + get_block(in_channel=64, depth=64, num_units=3), + get_block(in_channel=64, depth=128, num_units=13), + get_block(in_channel=128, depth=256, num_units=30), + get_block(in_channel=256, depth=512, num_units=3) + ] + elif num_layers == 152: + blocks = [ + get_block(in_channel=64, depth=64, num_units=3), + get_block(in_channel=64, depth=128, num_units=8), + get_block(in_channel=128, depth=256, num_units=36), + get_block(in_channel=256, depth=512, num_units=3) + ] + + return blocks + + +class Backbone(Module): + + def __init__(self, input_size, num_layers, mode='ir'): + super(Backbone, self).__init__() + assert input_size[0] in [ + 112, 224 + ], 'input_size should be [112, 112] or [224, 224]' + assert num_layers in [50, 100, + 152], 'num_layers should be 50, 100 or 152' + assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se' + blocks = get_blocks(num_layers) + if mode == 'ir': + unit_module = bottleneck_IR + elif mode == 'ir_se': + unit_module = bottleneck_IR_SE + self.input_layer = Sequential( + Conv2d(3, 64, (3, 3), 1, 1, bias=False), BatchNorm2d(64), + PReLU(64)) + if input_size[0] == 112: + self.output_layer = Sequential( + BatchNorm2d(512), Dropout(0.4), Flatten(), + Linear(512 * 7 * 7, 512), BatchNorm1d(512, affine=False)) + else: + self.output_layer = Sequential( + BatchNorm2d(512), Dropout(0.4), Flatten(), + Linear(512 * 14 * 14, 512), BatchNorm1d(512, affine=False)) + + modules = [] + for block in blocks: + for bottleneck in block: + modules.append( + unit_module(bottleneck.in_channel, bottleneck.depth, + bottleneck.stride)) + self.body = Sequential(*modules) + + self._initialize_weights() + + def forward(self, x): + x = self.input_layer(x) + x = self.body(x) + conv_out = x.view(x.shape[0], -1) + x = self.output_layer(x) + + return l2_norm(x), conv_out + + def _initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_( + m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + m.bias.data.zero_() + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + elif isinstance(m, nn.Linear): + nn.init.kaiming_normal_( + m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + m.bias.data.zero_() + + +def IR_50(input_size): + """Constructs a ir-50 model. + """ + model = Backbone(input_size, 50, 'ir') + + return model + + +def IR_101(input_size): + """Constructs a ir-101 model. + """ + model = Backbone(input_size, 100, 'ir') + + return model + + +def IR_152(input_size): + """Constructs a ir-152 model. + """ + model = Backbone(input_size, 152, 'ir') + + return model + + +def IR_SE_50(input_size): + """Constructs a ir_se-50 model. + """ + model = Backbone(input_size, 50, 'ir_se') + + return model + + +def IR_SE_101(input_size): + """Constructs a ir_se-101 model. + """ + model = Backbone(input_size, 100, 'ir_se') + + return model + + +def IR_SE_152(input_size): + """Constructs a ir_se-152 model. + """ + model = Backbone(input_size, 152, 'ir_se') + + return model diff --git a/modelscope/models/cv/image_face_fusion/network/ops.py b/modelscope/models/cv/image_face_fusion/network/ops.py new file mode 100644 index 00000000..d168e691 --- /dev/null +++ b/modelscope/models/cv/image_face_fusion/network/ops.py @@ -0,0 +1,211 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn import Parameter + + +def init_func(m, init_type='xavier', gain=0.02): + classname = m.__class__.__name__ + if classname.find('BatchNorm2d') != -1: + if hasattr(m, 'weight') and m.weight is not None: + nn.init.normal_(m.weight, 1.0, gain) + if hasattr(m, 'bias') and m.bias is not None: + nn.init.constant_(m.bias, 0.0) + elif hasattr(m, 'weight') and (classname.find('Conv') != -1 + or classname.find('Linear') != -1): + if init_type == 'normal': + nn.init.normal_(m.weight, 0.0, gain) + elif init_type == 'xavier': + nn.init.xavier_normal_(m.weight, gain=gain) + elif init_type == 'xavier_uniform': + nn.init.xavier_uniform_(m.weight, gain=1.0) + elif init_type == 'kaiming': + nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in') + elif init_type == 'orthogonal': + nn.init.orthogonal_(m.weight, gain=gain) + elif init_type == 'none': # uses pytorch's default init method + m.reset_parameters() + else: + raise NotImplementedError( + 'initialization method [%s] is not implemented' % init_type) + if hasattr(m, 'bias') and m.bias is not None: + nn.init.constant_(m.bias, 0.0) + elif hasattr(m, 'weight_bar') and (classname.find('Conv') != -1): + if init_type == 'normal': + nn.init.normal_(m.weight_bar, 0.0, gain) + elif init_type == 'xavier': + nn.init.xavier_normal_(m.weight_bar, gain=gain) + elif init_type == 'xavier_uniform': + nn.init.xavier_uniform_(m.weight_bar, gain=1.0) + elif init_type == 'kaiming': + nn.init.kaiming_normal_(m.weight_bar, a=0, mode='fan_in') + elif init_type == 'orthogonal': + nn.init.orthogonal_(m.weight_bar, gain=gain) + elif init_type == 'none': # uses pytorch's default init method + m.reset_parameters() + else: + raise NotImplementedError( + 'initialization method [%s] is not implemented' % init_type) + if hasattr(m, 'bias') and m.bias is not None: + nn.init.constant_(m.bias, 0.0) + + +def l2normalize(v, eps=1e-12): + return v / (v.norm() + eps) + + +class SpectralNorm(nn.Module): + + def __init__(self, module, name='weight', power_iterations=1): + super(SpectralNorm, self).__init__() + self.module = module + self.name = name + self.power_iterations = power_iterations + self._make_params() + + def _update_u_v(self): + u = getattr(self.module, self.name + '_u') + v = getattr(self.module, self.name + '_v') + w = getattr(self.module, self.name + '_bar') + + height = w.data.shape[0] + for _ in range(self.power_iterations): + v.data = l2normalize( + torch.mv(torch.t(w.view(height, -1).data), u.data)) + u.data = l2normalize(torch.mv(w.view(height, -1).data, v.data)) + + sigma = u.dot(w.view(height, -1).mv(v)) + setattr(self.module, self.name, w / sigma.expand_as(w)) + + def _noupdate_u_v(self): + u = getattr(self.module, self.name + '_u') + v = getattr(self.module, self.name + '_v') + w = getattr(self.module, self.name + '_bar') + + height = w.data.shape[0] + sigma = u.dot(w.view(height, -1).mv(v)) + setattr(self.module, self.name, w / sigma.expand_as(w)) + + def _make_params(self): + w = getattr(self.module, self.name) + + height = w.data.shape[0] + width = w.view(height, -1).data.shape[1] + + u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False) + v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False) + u.data = l2normalize(u.data) + v.data = l2normalize(v.data) + w_bar = Parameter(w.data) + + del self.module._parameters[self.name] + + self.module.register_parameter(self.name + '_u', u) + self.module.register_parameter(self.name + '_v', v) + self.module.register_parameter(self.name + '_bar', w_bar) + + def forward(self, *args): + if self.module.training: + self._update_u_v() + else: + self._noupdate_u_v() + return self.module.forward(*args) + + +def convert_affinematrix_to_homography(A): + H = torch.nn.functional.pad(A, [0, 0, 0, 1], 'constant', value=0.0) + H[..., -1, -1] += 1.0 + return H + + +def normal_transform_pixel(height, width, eps=1e-14): + tr_mat = torch.tensor([[1.0, 0.0, -1.0], [0.0, 1.0, -1.0], [0.0, 0.0, + 1.0]]) # 3x3 + + # prevent divide by zero bugs + width_denom = eps if width == 1 else width - 1.0 + height_denom = eps if height == 1 else height - 1.0 + + tr_mat[0, 0] = tr_mat[0, 0] * 2.0 / width_denom + tr_mat[1, 1] = tr_mat[1, 1] * 2.0 / height_denom + + return tr_mat.unsqueeze(0) # 1x3x3 + + +def _torch_inverse_cast(input): + if not isinstance(input, torch.Tensor): + raise AssertionError( + f'Input must be torch.Tensor. Got: {type(input)}.') + dtype = input.dtype + if dtype not in (torch.float32, torch.float64): + dtype = torch.float32 + return torch.inverse(input.to(dtype)).to(input.dtype) + + +def normalize_homography(dst_pix_trans_src_pix, dsize_src, dsize_dst): + if not isinstance(dst_pix_trans_src_pix, torch.Tensor): + raise TypeError( + f'Input type is not a torch.Tensor. Got {type(dst_pix_trans_src_pix)}' + ) + + if not (len(dst_pix_trans_src_pix.shape) == 3 + or dst_pix_trans_src_pix.shape[-2:] == (3, 3)): + raise ValueError( + f'Input dst_pix_trans_src_pix must be a Bx3x3 tensor. Got {dst_pix_trans_src_pix.shape}' + ) + + # source and destination sizes + src_h, src_w = dsize_src + dst_h, dst_w = dsize_dst + + # compute the transformation pixel/norm for src/dst + src_norm_trans_src_pix: torch.Tensor = normal_transform_pixel( + src_h, src_w).to(dst_pix_trans_src_pix) + + src_pix_trans_src_norm = _torch_inverse_cast(src_norm_trans_src_pix) + dst_norm_trans_dst_pix = normal_transform_pixel( + dst_h, dst_w).to(dst_pix_trans_src_pix) + + # compute chain transformations + dst_norm_trans_src_norm = dst_norm_trans_dst_pix @ ( + dst_pix_trans_src_pix @ src_pix_trans_src_norm) + return dst_norm_trans_src_norm + + +def warp_affine_torch(src, + M, + dsize, + mode='bilinear', + padding_mode='zeros', + align_corners=True): + + if not isinstance(src, torch.Tensor): + raise TypeError( + f'Input src type is not a torch.Tensor. Got {type(src)}') + + if not isinstance(M, torch.Tensor): + raise TypeError(f'Input M type is not a torch.Tensor. Got {type(M)}') + + if not len(src.shape) == 4: + raise ValueError( + f'Input src must be a BxCxHxW tensor. Got {src.shape}') + + if not (len(M.shape) == 3 or M.shape[-2:] == (2, 3)): + raise ValueError(f'Input M must be a Bx2x3 tensor. Got {M.shape}') + + B, C, H, W = src.size() + + # we generate a 3x3 transformation matrix from 2x3 affine + M_3x3 = convert_affinematrix_to_homography(M) + dst_norm_trans_src_norm = normalize_homography(M_3x3, (H, W), dsize) + src_norm_trans_dst_norm = _torch_inverse_cast(dst_norm_trans_src_norm) + grid = F.affine_grid( + src_norm_trans_dst_norm[:, :2, :], [B, C, dsize[0], dsize[1]], + align_corners=align_corners) + return F.grid_sample( + src, + grid, + align_corners=align_corners, + mode=mode, + padding_mode=padding_mode) diff --git a/modelscope/pipelines/builder.py b/modelscope/pipelines/builder.py index 31da8134..f3629469 100644 --- a/modelscope/pipelines/builder.py +++ b/modelscope/pipelines/builder.py @@ -256,6 +256,10 @@ DEFAULT_MODEL_FOR_PIPELINE = { Tasks.image_multi_view_depth_estimation: ( Pipelines.image_multi_view_depth_estimation, 'damo/cv_casmvs_multi-view-depth-estimation_general'), + Tasks.image_body_reshaping: (Pipelines.image_body_reshaping, + 'damo/cv_flow-based-body-reshaping_damo'), + Tasks.image_face_fusion: (Pipelines.image_face_fusion, + 'damo/cv_unet-image-face-fusion_damo'), } diff --git a/modelscope/pipelines/cv/image_face_fusion_pipeline.py b/modelscope/pipelines/cv/image_face_fusion_pipeline.py new file mode 100644 index 00000000..3ba253e1 --- /dev/null +++ b/modelscope/pipelines/cv/image_face_fusion_pipeline.py @@ -0,0 +1,68 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import Any, Dict + +import numpy as np + +from modelscope.metainfo import Pipelines +from modelscope.outputs import OutputKeys +from modelscope.pipelines.base import Input, Pipeline +from modelscope.pipelines.builder import PIPELINES +from modelscope.preprocessors import LoadImage +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.logger import get_logger + +logger = get_logger() + + +@PIPELINES.register_module( + Tasks.image_face_fusion, module_name=Pipelines.image_face_fusion) +class ImageFaceFusionPipeline(Pipeline): + """ Image face fusion pipeline + Example: + + python + >>> from modelscope.pipelines import pipeline + >>> image_face_fusion = pipeline(Tasks.image_face_fusion, + model='damo/cv_unet-image-face-fusion_damo') + >>> image_face_fusion({ + 'template': 'facefusion_template.jpg', # template path (str) + 'image': 'facefusion_user.jpg', # user path (str) + }) + { + "output_img": [H * W * 3] 0~255, we can use cv2.imwrite to save output_img as an image. + } + >>> # + """ + + def __init__(self, model: str, **kwargs): + """ + use `model` to create image-face-fusion pipeline for prediction + Args: + model: model id on modelscope hub. + """ + super().__init__(model=model, **kwargs) + logger.info('image face fusion model init done') + + def preprocess(self, + template: Input, + user: Input = None) -> Dict[str, Any]: + if type(template) is dict: # for demo service + user = template['user'] + template = template['template'] + + template_img = LoadImage.convert_to_ndarray(template) + user_img = LoadImage.convert_to_ndarray(user) + + result = {'template': template_img, 'user': user_img} + return result + + def forward(self, input: Dict[str, Any]) -> Dict[str, Any]: + template_img = input['template'] + user_img = input['user'] + output = self.model.inference(template_img, user_img) + result = {'outputs': output} + return result + + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + output_img = inputs['outputs'] + return {OutputKeys.OUTPUT_IMG: output_img} diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index c95b5502..6bb9a142 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -78,7 +78,7 @@ class CVTasks(object): image_portrait_stylization = 'image-portrait-stylization' image_body_reshaping = 'image-body-reshaping' image_embedding = 'image-embedding' - + image_face_fusion = 'image-face-fusion' product_retrieval_embedding = 'product-retrieval-embedding' # video recognition diff --git a/tests/pipelines/test_image_face_fusion.py b/tests/pipelines/test_image_face_fusion.py new file mode 100644 index 00000000..fde15edf --- /dev/null +++ b/tests/pipelines/test_image_face_fusion.py @@ -0,0 +1,59 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import unittest + +import cv2 + +from modelscope.hub.snapshot_download import snapshot_download +from modelscope.outputs import OutputKeys +from modelscope.pipelines import pipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.demo_utils import DemoCompatibilityCheck +from modelscope.utils.test_utils import test_level + + +class ImageFaceFusionTest(unittest.TestCase, DemoCompatibilityCheck): + + def setUp(self) -> None: + self.task = Tasks.image_face_fusion + self.model_id = 'damo/cv_unet-image-face-fusion_damo' + self.template_img = 'data/test/images/facefusion_template.jpg' + self.user_img = 'data/test/images/facefusion_user.jpg' + + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') + def test_run_by_direct_model_download(self): + snapshot_path = snapshot_download(self.model_id) + print('snapshot_path: {}'.format(snapshot_path)) + image_face_fusion = pipeline( + Tasks.image_face_fusion, model=snapshot_path) + + result = image_face_fusion( + dict(template=self.template_img, user=self.user_img)) + cv2.imwrite('result_facefusion.png', result[OutputKeys.OUTPUT_IMG]) + print('facefusion.test_run_direct_model_download done') + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_modelhub(self): + image_face_fusion = pipeline( + Tasks.image_face_fusion, model=self.model_id) + + result = image_face_fusion( + dict(template=self.template_img, user=self.user_img)) + cv2.imwrite('result_facefusion.png', result[OutputKeys.OUTPUT_IMG]) + print('facefusion.test_run_modelhub done') + + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') + def test_run_modelhub_default_model(self): + image_face_fusion = pipeline(Tasks.image_face_fusion) + + result = image_face_fusion( + dict(template=self.template_img, user=self.user_img)) + cv2.imwrite('result_facefusion.png', result[OutputKeys.OUTPUT_IMG]) + print('facefusion.test_run_modelhub_default_model done') + + @unittest.skip('demo compatibility test is only enabled on a needed-basis') + def test_demo_compatibility(self): + self.compatibility_check() + + +if __name__ == '__main__': + unittest.main()