Init

CODE_OF_CONDUCT.md

@@ -0,0 +1,9 @@
+# Microsoft Open Source Code of Conduct
+
+This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
+
+Resources:
+
+- [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/)
+- [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/)
+- Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns

Dockerfile

@@ -0,0 +1,43 @@
+FROM nvidia/cuda:11.1-base-ubuntu20.04
+
+RUN apt update && DEBIAN_FRONTEND=noninteractive apt install git bzip2 wget unzip python3-pip python3-dev cmake libgl1-mesa-dev python-is-python3 libgtk2.0-dev -yq
+ADD . /app
+WORKDIR /app
+RUN cd Face_Enhancement/models/networks/ &&\
+  git clone https://github.com/vacancy/Synchronized-BatchNorm-PyTorch &&\
+  cp -rf Synchronized-BatchNorm-PyTorch/sync_batchnorm . &&\
+  cd ../../../
+
+RUN cd Global/detection_models &&\
+  git clone https://github.com/vacancy/Synchronized-BatchNorm-PyTorch &&\
+  cp -rf Synchronized-BatchNorm-PyTorch/sync_batchnorm . &&\
+  cd ../../
+
+RUN cd Face_Detection/ &&\
+  wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2 &&\
+  bzip2 -d shape_predictor_68_face_landmarks.dat.bz2 &&\
+  cd ../ 
+
+RUN cd Face_Enhancement/ &&\
+  wget https://facevc.blob.core.windows.net/zhanbo/old_photo/pretrain/Face_Enhancement/checkpoints.zip &&\
+  unzip checkpoints.zip &&\
+  cd ../ &&\
+  cd Global/ &&\
+  wget https://facevc.blob.core.windows.net/zhanbo/old_photo/pretrain/Global/checkpoints.zip &&\
+  unzip checkpoints.zip &&\
+  rm -f checkpoints.zip &&\
+  cd ../
+
+RUN pip3 install numpy
+
+RUN pip3 install dlib
+
+RUN pip3 install -r requirements.txt
+
+RUN git clone https://github.com/NVlabs/SPADE.git
+
+RUN cd SPADE/ && pip3 install -r requirements.txt
+
+RUN cd ..
+
+CMD ["python3", "run.py"]

Face_Detection/align_warp_back_multiple_dlib.py

@@ -0,0 +1,437 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+import numpy as np
+import skimage.io as io
+
+# from face_sdk import FaceDetection
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+from skimage.transform import SimilarityTransform
+from skimage.transform import warp
+from PIL import Image, ImageFilter
+import torch.nn.functional as F
+import torchvision as tv
+import torchvision.utils as vutils
+import time
+import cv2
+import os
+from skimage import img_as_ubyte
+import json
+import argparse
+import dlib
+
+
+def calculate_cdf(histogram):
+    """
+    This method calculates the cumulative distribution function
+    :param array histogram: The values of the histogram
+    :return: normalized_cdf: The normalized cumulative distribution function
+    :rtype: array
+    """
+    # Get the cumulative sum of the elements
+    cdf = histogram.cumsum()
+
+    # Normalize the cdf
+    normalized_cdf = cdf / float(cdf.max())
+
+    return normalized_cdf
+
+
+def calculate_lookup(src_cdf, ref_cdf):
+    """
+    This method creates the lookup table
+    :param array src_cdf: The cdf for the source image
+    :param array ref_cdf: The cdf for the reference image
+    :return: lookup_table: The lookup table
+    :rtype: array
+    """
+    lookup_table = np.zeros(256)
+    lookup_val = 0
+    for src_pixel_val in range(len(src_cdf)):
+        lookup_val
+        for ref_pixel_val in range(len(ref_cdf)):
+            if ref_cdf[ref_pixel_val] >= src_cdf[src_pixel_val]:
+                lookup_val = ref_pixel_val
+                break
+        lookup_table[src_pixel_val] = lookup_val
+    return lookup_table
+
+
+def match_histograms(src_image, ref_image):
+    """
+    This method matches the source image histogram to the
+    reference signal
+    :param image src_image: The original source image
+    :param image  ref_image: The reference image
+    :return: image_after_matching
+    :rtype: image (array)
+    """
+    # Split the images into the different color channels
+    # b means blue, g means green and r means red
+    src_b, src_g, src_r = cv2.split(src_image)
+    ref_b, ref_g, ref_r = cv2.split(ref_image)
+
+    # Compute the b, g, and r histograms separately
+    # The flatten() Numpy method returns a copy of the array c
+    # collapsed into one dimension.
+    src_hist_blue, bin_0 = np.histogram(src_b.flatten(), 256, [0, 256])
+    src_hist_green, bin_1 = np.histogram(src_g.flatten(), 256, [0, 256])
+    src_hist_red, bin_2 = np.histogram(src_r.flatten(), 256, [0, 256])
+    ref_hist_blue, bin_3 = np.histogram(ref_b.flatten(), 256, [0, 256])
+    ref_hist_green, bin_4 = np.histogram(ref_g.flatten(), 256, [0, 256])
+    ref_hist_red, bin_5 = np.histogram(ref_r.flatten(), 256, [0, 256])
+
+    # Compute the normalized cdf for the source and reference image
+    src_cdf_blue = calculate_cdf(src_hist_blue)
+    src_cdf_green = calculate_cdf(src_hist_green)
+    src_cdf_red = calculate_cdf(src_hist_red)
+    ref_cdf_blue = calculate_cdf(ref_hist_blue)
+    ref_cdf_green = calculate_cdf(ref_hist_green)
+    ref_cdf_red = calculate_cdf(ref_hist_red)
+
+    # Make a separate lookup table for each color
+    blue_lookup_table = calculate_lookup(src_cdf_blue, ref_cdf_blue)
+    green_lookup_table = calculate_lookup(src_cdf_green, ref_cdf_green)
+    red_lookup_table = calculate_lookup(src_cdf_red, ref_cdf_red)
+
+    # Use the lookup function to transform the colors of the original
+    # source image
+    blue_after_transform = cv2.LUT(src_b, blue_lookup_table)
+    green_after_transform = cv2.LUT(src_g, green_lookup_table)
+    red_after_transform = cv2.LUT(src_r, red_lookup_table)
+
+    # Put the image back together
+    image_after_matching = cv2.merge([blue_after_transform, green_after_transform, red_after_transform])
+    image_after_matching = cv2.convertScaleAbs(image_after_matching)
+
+    return image_after_matching
+
+
+def _standard_face_pts():
+    pts = (
+        np.array([196.0, 226.0, 316.0, 226.0, 256.0, 286.0, 220.0, 360.4, 292.0, 360.4], np.float32) / 256.0
+        - 1.0
+    )
+
+    return np.reshape(pts, (5, 2))
+
+
+def _origin_face_pts():
+    pts = np.array([196.0, 226.0, 316.0, 226.0, 256.0, 286.0, 220.0, 360.4, 292.0, 360.4], np.float32)
+
+    return np.reshape(pts, (5, 2))
+
+
+def compute_transformation_matrix(img, landmark, normalize, target_face_scale=1.0):
+
+    std_pts = _standard_face_pts()  # [-1,1]
+    target_pts = (std_pts * target_face_scale + 1) / 2 * 256.0
+
+    # print(target_pts)
+
+    h, w, c = img.shape
+    if normalize == True:
+        landmark[:, 0] = landmark[:, 0] / h * 2 - 1.0
+        landmark[:, 1] = landmark[:, 1] / w * 2 - 1.0
+
+    # print(landmark)
+
+    affine = SimilarityTransform()
+
+    affine.estimate(target_pts, landmark)
+
+    return affine
+
+
+def compute_inverse_transformation_matrix(img, landmark, normalize, target_face_scale=1.0):
+
+    std_pts = _standard_face_pts()  # [-1,1]
+    target_pts = (std_pts * target_face_scale + 1) / 2 * 256.0
+
+    # print(target_pts)
+
+    h, w, c = img.shape
+    if normalize == True:
+        landmark[:, 0] = landmark[:, 0] / h * 2 - 1.0
+        landmark[:, 1] = landmark[:, 1] / w * 2 - 1.0
+
+    # print(landmark)
+
+    affine = SimilarityTransform()
+
+    affine.estimate(landmark, target_pts)
+
+    return affine
+
+
+def show_detection(image, box, landmark):
+    plt.imshow(image)
+    print(box[2] - box[0])
+    plt.gca().add_patch(
+        Rectangle(
+            (box[1], box[0]), box[2] - box[0], box[3] - box[1], linewidth=1, edgecolor="r", facecolor="none"
+        )
+    )
+    plt.scatter(landmark[0][0], landmark[0][1])
+    plt.scatter(landmark[1][0], landmark[1][1])
+    plt.scatter(landmark[2][0], landmark[2][1])
+    plt.scatter(landmark[3][0], landmark[3][1])
+    plt.scatter(landmark[4][0], landmark[4][1])
+    plt.show()
+
+
+def affine2theta(affine, input_w, input_h, target_w, target_h):
+    # param = np.linalg.inv(affine)
+    param = affine
+    theta = np.zeros([2, 3])
+    theta[0, 0] = param[0, 0] * input_h / target_h
+    theta[0, 1] = param[0, 1] * input_w / target_h
+    theta[0, 2] = (2 * param[0, 2] + param[0, 0] * input_h + param[0, 1] * input_w) / target_h - 1
+    theta[1, 0] = param[1, 0] * input_h / target_w
+    theta[1, 1] = param[1, 1] * input_w / target_w
+    theta[1, 2] = (2 * param[1, 2] + param[1, 0] * input_h + param[1, 1] * input_w) / target_w - 1
+    return theta
+
+
+def blur_blending(im1, im2, mask):
+
+    mask *= 255.0
+
+    kernel = np.ones((10, 10), np.uint8)
+    mask = cv2.erode(mask, kernel, iterations=1)
+
+    mask = Image.fromarray(mask.astype("uint8")).convert("L")
+    im1 = Image.fromarray(im1.astype("uint8"))
+    im2 = Image.fromarray(im2.astype("uint8"))
+
+    mask_blur = mask.filter(ImageFilter.GaussianBlur(20))
+    im = Image.composite(im1, im2, mask)
+
+    im = Image.composite(im, im2, mask_blur)
+
+    return np.array(im) / 255.0
+
+
+def blur_blending_cv2(im1, im2, mask):
+
+    mask *= 255.0
+
+    kernel = np.ones((9, 9), np.uint8)
+    mask = cv2.erode(mask, kernel, iterations=3)
+
+    mask_blur = cv2.GaussianBlur(mask, (25, 25), 0)
+    mask_blur /= 255.0
+
+    im = im1 * mask_blur + (1 - mask_blur) * im2
+
+    im /= 255.0
+    im = np.clip(im, 0.0, 1.0)
+
+    return im
+
+
+# def Poisson_blending(im1,im2,mask):
+
+
+#     Image.composite(
+def Poisson_blending(im1, im2, mask):
+
+    # mask=1-mask
+    mask *= 255
+    kernel = np.ones((10, 10), np.uint8)
+    mask = cv2.erode(mask, kernel, iterations=1)
+    mask /= 255
+    mask = 1 - mask
+    mask *= 255
+
+    mask = mask[:, :, 0]
+    width, height, channels = im1.shape
+    center = (int(height / 2), int(width / 2))
+    result = cv2.seamlessClone(
+        im2.astype("uint8"), im1.astype("uint8"), mask.astype("uint8"), center, cv2.MIXED_CLONE
+    )
+
+    return result / 255.0
+
+
+def Poisson_B(im1, im2, mask, center):
+
+    mask *= 255
+
+    result = cv2.seamlessClone(
+        im2.astype("uint8"), im1.astype("uint8"), mask.astype("uint8"), center, cv2.NORMAL_CLONE
+    )
+
+    return result / 255
+
+
+def seamless_clone(old_face, new_face, raw_mask):
+
+    height, width, _ = old_face.shape
+    height = height // 2
+    width = width // 2
+
+    y_indices, x_indices, _ = np.nonzero(raw_mask)
+    y_crop = slice(np.min(y_indices), np.max(y_indices))
+    x_crop = slice(np.min(x_indices), np.max(x_indices))
+    y_center = int(np.rint((np.max(y_indices) + np.min(y_indices)) / 2 + height))
+    x_center = int(np.rint((np.max(x_indices) + np.min(x_indices)) / 2 + width))
+
+    insertion = np.rint(new_face[y_crop, x_crop] * 255.0).astype("uint8")
+    insertion_mask = np.rint(raw_mask[y_crop, x_crop] * 255.0).astype("uint8")
+    insertion_mask[insertion_mask != 0] = 255
+    prior = np.rint(np.pad(old_face * 255.0, ((height, height), (width, width), (0, 0)), "constant")).astype(
+        "uint8"
+    )
+    # if np.sum(insertion_mask) == 0:
+    n_mask = insertion_mask[1:-1, 1:-1, :]
+    n_mask = cv2.copyMakeBorder(n_mask, 1, 1, 1, 1, cv2.BORDER_CONSTANT, 0)
+    print(n_mask.shape)
+    x, y, w, h = cv2.boundingRect(n_mask[:, :, 0])
+    if w < 4 or h < 4:
+        blended = prior
+    else:
+        blended = cv2.seamlessClone(
+            insertion,  # pylint: disable=no-member
+            prior,
+            insertion_mask,
+            (x_center, y_center),
+            cv2.NORMAL_CLONE,
+        )  # pylint: disable=no-member
+
+    blended = blended[height:-height, width:-width]
+
+    return blended.astype("float32") / 255.0
+
+
+def get_landmark(face_landmarks, id):
+    part = face_landmarks.part(id)
+    x = part.x
+    y = part.y
+
+    return (x, y)
+
+
+def search(face_landmarks):
+
+    x1, y1 = get_landmark(face_landmarks, 36)
+    x2, y2 = get_landmark(face_landmarks, 39)
+    x3, y3 = get_landmark(face_landmarks, 42)
+    x4, y4 = get_landmark(face_landmarks, 45)
+
+    x_nose, y_nose = get_landmark(face_landmarks, 30)
+
+    x_left_mouth, y_left_mouth = get_landmark(face_landmarks, 48)
+    x_right_mouth, y_right_mouth = get_landmark(face_landmarks, 54)
+
+    x_left_eye = int((x1 + x2) / 2)
+    y_left_eye = int((y1 + y2) / 2)
+    x_right_eye = int((x3 + x4) / 2)
+    y_right_eye = int((y3 + y4) / 2)
+
+    results = np.array(
+        [
+            [x_left_eye, y_left_eye],
+            [x_right_eye, y_right_eye],
+            [x_nose, y_nose],
+            [x_left_mouth, y_left_mouth],
+            [x_right_mouth, y_right_mouth],
+        ]
+    )
+
+    return results
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--origin_url", type=str, default="./", help="origin images")
+    parser.add_argument("--replace_url", type=str, default="./", help="restored faces")
+    parser.add_argument("--save_url", type=str, default="./save")
+    opts = parser.parse_args()
+
+    origin_url = opts.origin_url
+    replace_url = opts.replace_url
+    save_url = opts.save_url
+
+    if not os.path.exists(save_url):
+        os.makedirs(save_url)
+
+    face_detector = dlib.get_frontal_face_detector()
+    landmark_locator = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
+
+    count = 0
+
+    for x in os.listdir(origin_url):
+        img_url = os.path.join(origin_url, x)
+        pil_img = Image.open(img_url).convert("RGB")
+
+        origin_width, origin_height = pil_img.size
+        image = np.array(pil_img)
+
+        start = time.time()
+        faces = face_detector(image)
+        done = time.time()
+
+        if len(faces) == 0:
+            print("Warning: There is no face in %s" % (x))
+            continue
+
+        blended = image
+        for face_id in range(len(faces)):
+
+            current_face = faces[face_id]
+            face_landmarks = landmark_locator(image, current_face)
+            current_fl = search(face_landmarks)
+
+            forward_mask = np.ones_like(image).astype("uint8")
+            affine = compute_transformation_matrix(image, current_fl, False, target_face_scale=1.3)
+            aligned_face = warp(image, affine, output_shape=(256, 256, 3), preserve_range=True)
+            forward_mask = warp(
+                forward_mask, affine, output_shape=(256, 256, 3), order=0, preserve_range=True
+            )
+
+            affine_inverse = affine.inverse
+            cur_face = aligned_face
+            if replace_url != "":
+
+                face_name = x[:-4] + "_" + str(face_id + 1) + ".png"
+                cur_url = os.path.join(replace_url, face_name)
+                restored_face = Image.open(cur_url).convert("RGB")
+                restored_face = np.array(restored_face)
+                cur_face = restored_face
+
+            ## Histogram Color matching
+            A = cv2.cvtColor(aligned_face.astype("uint8"), cv2.COLOR_RGB2BGR)
+            B = cv2.cvtColor(cur_face.astype("uint8"), cv2.COLOR_RGB2BGR)
+            B = match_histograms(B, A)
+            cur_face = cv2.cvtColor(B.astype("uint8"), cv2.COLOR_BGR2RGB)
+
+            warped_back = warp(
+                cur_face,
+                affine_inverse,
+                output_shape=(origin_height, origin_width, 3),
+                order=3,
+                preserve_range=True,
+            )
+
+            backward_mask = warp(
+                forward_mask,
+                affine_inverse,
+                output_shape=(origin_height, origin_width, 3),
+                order=0,
+                preserve_range=True,
+            )  ## Nearest neighbour
+
+            blended = blur_blending_cv2(warped_back, blended, backward_mask)
+            blended *= 255.0
+
+        io.imsave(os.path.join(save_url, x), img_as_ubyte(blended / 255.0))
+
+        count += 1
+
+        if count % 1000 == 0:
+            print("%d have finished ..." % (count))
+

Face_Detection/align_warp_back_multiple_dlib_HR.py

@@ -0,0 +1,437 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+import numpy as np
+import skimage.io as io
+
+# from face_sdk import FaceDetection
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+from skimage.transform import SimilarityTransform
+from skimage.transform import warp
+from PIL import Image, ImageFilter
+import torch.nn.functional as F
+import torchvision as tv
+import torchvision.utils as vutils
+import time
+import cv2
+import os
+from skimage import img_as_ubyte
+import json
+import argparse
+import dlib
+
+
+def calculate_cdf(histogram):
+    """
+    This method calculates the cumulative distribution function
+    :param array histogram: The values of the histogram
+    :return: normalized_cdf: The normalized cumulative distribution function
+    :rtype: array
+    """
+    # Get the cumulative sum of the elements
+    cdf = histogram.cumsum()
+
+    # Normalize the cdf
+    normalized_cdf = cdf / float(cdf.max())
+
+    return normalized_cdf
+
+
+def calculate_lookup(src_cdf, ref_cdf):
+    """
+    This method creates the lookup table
+    :param array src_cdf: The cdf for the source image
+    :param array ref_cdf: The cdf for the reference image
+    :return: lookup_table: The lookup table
+    :rtype: array
+    """
+    lookup_table = np.zeros(256)
+    lookup_val = 0
+    for src_pixel_val in range(len(src_cdf)):
+        lookup_val
+        for ref_pixel_val in range(len(ref_cdf)):
+            if ref_cdf[ref_pixel_val] >= src_cdf[src_pixel_val]:
+                lookup_val = ref_pixel_val
+                break
+        lookup_table[src_pixel_val] = lookup_val
+    return lookup_table
+
+
+def match_histograms(src_image, ref_image):
+    """
+    This method matches the source image histogram to the
+    reference signal
+    :param image src_image: The original source image
+    :param image  ref_image: The reference image
+    :return: image_after_matching
+    :rtype: image (array)
+    """
+    # Split the images into the different color channels
+    # b means blue, g means green and r means red
+    src_b, src_g, src_r = cv2.split(src_image)
+    ref_b, ref_g, ref_r = cv2.split(ref_image)
+
+    # Compute the b, g, and r histograms separately
+    # The flatten() Numpy method returns a copy of the array c
+    # collapsed into one dimension.
+    src_hist_blue, bin_0 = np.histogram(src_b.flatten(), 256, [0, 256])
+    src_hist_green, bin_1 = np.histogram(src_g.flatten(), 256, [0, 256])
+    src_hist_red, bin_2 = np.histogram(src_r.flatten(), 256, [0, 256])
+    ref_hist_blue, bin_3 = np.histogram(ref_b.flatten(), 256, [0, 256])
+    ref_hist_green, bin_4 = np.histogram(ref_g.flatten(), 256, [0, 256])
+    ref_hist_red, bin_5 = np.histogram(ref_r.flatten(), 256, [0, 256])
+
+    # Compute the normalized cdf for the source and reference image
+    src_cdf_blue = calculate_cdf(src_hist_blue)
+    src_cdf_green = calculate_cdf(src_hist_green)
+    src_cdf_red = calculate_cdf(src_hist_red)
+    ref_cdf_blue = calculate_cdf(ref_hist_blue)
+    ref_cdf_green = calculate_cdf(ref_hist_green)
+    ref_cdf_red = calculate_cdf(ref_hist_red)
+
+    # Make a separate lookup table for each color
+    blue_lookup_table = calculate_lookup(src_cdf_blue, ref_cdf_blue)
+    green_lookup_table = calculate_lookup(src_cdf_green, ref_cdf_green)
+    red_lookup_table = calculate_lookup(src_cdf_red, ref_cdf_red)
+
+    # Use the lookup function to transform the colors of the original
+    # source image
+    blue_after_transform = cv2.LUT(src_b, blue_lookup_table)
+    green_after_transform = cv2.LUT(src_g, green_lookup_table)
+    red_after_transform = cv2.LUT(src_r, red_lookup_table)
+
+    # Put the image back together
+    image_after_matching = cv2.merge([blue_after_transform, green_after_transform, red_after_transform])
+    image_after_matching = cv2.convertScaleAbs(image_after_matching)
+
+    return image_after_matching
+
+
+def _standard_face_pts():
+    pts = (
+        np.array([196.0, 226.0, 316.0, 226.0, 256.0, 286.0, 220.0, 360.4, 292.0, 360.4], np.float32) / 256.0
+        - 1.0
+    )
+
+    return np.reshape(pts, (5, 2))
+
+
+def _origin_face_pts():
+    pts = np.array([196.0, 226.0, 316.0, 226.0, 256.0, 286.0, 220.0, 360.4, 292.0, 360.4], np.float32)
+
+    return np.reshape(pts, (5, 2))
+
+
+def compute_transformation_matrix(img, landmark, normalize, target_face_scale=1.0):
+
+    std_pts = _standard_face_pts()  # [-1,1]
+    target_pts = (std_pts * target_face_scale + 1) / 2 * 512.0
+
+    # print(target_pts)
+
+    h, w, c = img.shape
+    if normalize == True:
+        landmark[:, 0] = landmark[:, 0] / h * 2 - 1.0
+        landmark[:, 1] = landmark[:, 1] / w * 2 - 1.0
+
+    # print(landmark)
+
+    affine = SimilarityTransform()
+
+    affine.estimate(target_pts, landmark)
+
+    return affine
+
+
+def compute_inverse_transformation_matrix(img, landmark, normalize, target_face_scale=1.0):
+
+    std_pts = _standard_face_pts()  # [-1,1]
+    target_pts = (std_pts * target_face_scale + 1) / 2 * 512.0
+
+    # print(target_pts)
+
+    h, w, c = img.shape
+    if normalize == True:
+        landmark[:, 0] = landmark[:, 0] / h * 2 - 1.0
+        landmark[:, 1] = landmark[:, 1] / w * 2 - 1.0
+
+    # print(landmark)
+
+    affine = SimilarityTransform()
+
+    affine.estimate(landmark, target_pts)
+
+    return affine
+
+
+def show_detection(image, box, landmark):
+    plt.imshow(image)
+    print(box[2] - box[0])
+    plt.gca().add_patch(
+        Rectangle(
+            (box[1], box[0]), box[2] - box[0], box[3] - box[1], linewidth=1, edgecolor="r", facecolor="none"
+        )
+    )
+    plt.scatter(landmark[0][0], landmark[0][1])
+    plt.scatter(landmark[1][0], landmark[1][1])
+    plt.scatter(landmark[2][0], landmark[2][1])
+    plt.scatter(landmark[3][0], landmark[3][1])
+    plt.scatter(landmark[4][0], landmark[4][1])
+    plt.show()
+
+
+def affine2theta(affine, input_w, input_h, target_w, target_h):
+    # param = np.linalg.inv(affine)
+    param = affine
+    theta = np.zeros([2, 3])
+    theta[0, 0] = param[0, 0] * input_h / target_h
+    theta[0, 1] = param[0, 1] * input_w / target_h
+    theta[0, 2] = (2 * param[0, 2] + param[0, 0] * input_h + param[0, 1] * input_w) / target_h - 1
+    theta[1, 0] = param[1, 0] * input_h / target_w
+    theta[1, 1] = param[1, 1] * input_w / target_w
+    theta[1, 2] = (2 * param[1, 2] + param[1, 0] * input_h + param[1, 1] * input_w) / target_w - 1
+    return theta
+
+
+def blur_blending(im1, im2, mask):
+
+    mask *= 255.0
+
+    kernel = np.ones((10, 10), np.uint8)
+    mask = cv2.erode(mask, kernel, iterations=1)
+
+    mask = Image.fromarray(mask.astype("uint8")).convert("L")
+    im1 = Image.fromarray(im1.astype("uint8"))
+    im2 = Image.fromarray(im2.astype("uint8"))
+
+    mask_blur = mask.filter(ImageFilter.GaussianBlur(20))
+    im = Image.composite(im1, im2, mask)
+
+    im = Image.composite(im, im2, mask_blur)
+
+    return np.array(im) / 255.0
+
+
+def blur_blending_cv2(im1, im2, mask):
+
+    mask *= 255.0
+
+    kernel = np.ones((9, 9), np.uint8)
+    mask = cv2.erode(mask, kernel, iterations=3)
+
+    mask_blur = cv2.GaussianBlur(mask, (25, 25), 0)
+    mask_blur /= 255.0
+
+    im = im1 * mask_blur + (1 - mask_blur) * im2
+
+    im /= 255.0
+    im = np.clip(im, 0.0, 1.0)
+
+    return im
+
+
+# def Poisson_blending(im1,im2,mask):
+
+
+#     Image.composite(
+def Poisson_blending(im1, im2, mask):
+
+    # mask=1-mask
+    mask *= 255
+    kernel = np.ones((10, 10), np.uint8)
+    mask = cv2.erode(mask, kernel, iterations=1)
+    mask /= 255
+    mask = 1 - mask
+    mask *= 255
+
+    mask = mask[:, :, 0]
+    width, height, channels = im1.shape
+    center = (int(height / 2), int(width / 2))
+    result = cv2.seamlessClone(
+        im2.astype("uint8"), im1.astype("uint8"), mask.astype("uint8"), center, cv2.MIXED_CLONE
+    )
+
+    return result / 255.0
+
+
+def Poisson_B(im1, im2, mask, center):
+
+    mask *= 255
+
+    result = cv2.seamlessClone(
+        im2.astype("uint8"), im1.astype("uint8"), mask.astype("uint8"), center, cv2.NORMAL_CLONE
+    )
+
+    return result / 255
+
+
+def seamless_clone(old_face, new_face, raw_mask):
+
+    height, width, _ = old_face.shape
+    height = height // 2
+    width = width // 2
+
+    y_indices, x_indices, _ = np.nonzero(raw_mask)
+    y_crop = slice(np.min(y_indices), np.max(y_indices))
+    x_crop = slice(np.min(x_indices), np.max(x_indices))
+    y_center = int(np.rint((np.max(y_indices) + np.min(y_indices)) / 2 + height))
+    x_center = int(np.rint((np.max(x_indices) + np.min(x_indices)) / 2 + width))
+
+    insertion = np.rint(new_face[y_crop, x_crop] * 255.0).astype("uint8")
+    insertion_mask = np.rint(raw_mask[y_crop, x_crop] * 255.0).astype("uint8")
+    insertion_mask[insertion_mask != 0] = 255
+    prior = np.rint(np.pad(old_face * 255.0, ((height, height), (width, width), (0, 0)), "constant")).astype(
+        "uint8"
+    )
+    # if np.sum(insertion_mask) == 0:
+    n_mask = insertion_mask[1:-1, 1:-1, :]
+    n_mask = cv2.copyMakeBorder(n_mask, 1, 1, 1, 1, cv2.BORDER_CONSTANT, 0)
+    print(n_mask.shape)
+    x, y, w, h = cv2.boundingRect(n_mask[:, :, 0])
+    if w < 4 or h < 4:
+        blended = prior
+    else:
+        blended = cv2.seamlessClone(
+            insertion,  # pylint: disable=no-member
+            prior,
+            insertion_mask,
+            (x_center, y_center),
+            cv2.NORMAL_CLONE,
+        )  # pylint: disable=no-member
+
+    blended = blended[height:-height, width:-width]
+
+    return blended.astype("float32") / 255.0
+
+
+def get_landmark(face_landmarks, id):
+    part = face_landmarks.part(id)
+    x = part.x
+    y = part.y
+
+    return (x, y)
+
+
+def search(face_landmarks):
+
+    x1, y1 = get_landmark(face_landmarks, 36)
+    x2, y2 = get_landmark(face_landmarks, 39)
+    x3, y3 = get_landmark(face_landmarks, 42)
+    x4, y4 = get_landmark(face_landmarks, 45)
+
+    x_nose, y_nose = get_landmark(face_landmarks, 30)
+
+    x_left_mouth, y_left_mouth = get_landmark(face_landmarks, 48)
+    x_right_mouth, y_right_mouth = get_landmark(face_landmarks, 54)
+
+    x_left_eye = int((x1 + x2) / 2)
+    y_left_eye = int((y1 + y2) / 2)
+    x_right_eye = int((x3 + x4) / 2)
+    y_right_eye = int((y3 + y4) / 2)
+
+    results = np.array(
+        [
+            [x_left_eye, y_left_eye],
+            [x_right_eye, y_right_eye],
+            [x_nose, y_nose],
+            [x_left_mouth, y_left_mouth],
+            [x_right_mouth, y_right_mouth],
+        ]
+    )
+
+    return results
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--origin_url", type=str, default="./", help="origin images")
+    parser.add_argument("--replace_url", type=str, default="./", help="restored faces")
+    parser.add_argument("--save_url", type=str, default="./save")
+    opts = parser.parse_args()
+
+    origin_url = opts.origin_url
+    replace_url = opts.replace_url
+    save_url = opts.save_url
+
+    if not os.path.exists(save_url):
+        os.makedirs(save_url)
+
+    face_detector = dlib.get_frontal_face_detector()
+    landmark_locator = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
+
+    count = 0
+
+    for x in os.listdir(origin_url):
+        img_url = os.path.join(origin_url, x)
+        pil_img = Image.open(img_url).convert("RGB")
+
+        origin_width, origin_height = pil_img.size
+        image = np.array(pil_img)
+
+        start = time.time()
+        faces = face_detector(image)
+        done = time.time()
+
+        if len(faces) == 0:
+            print("Warning: There is no face in %s" % (x))
+            continue
+
+        blended = image
+        for face_id in range(len(faces)):
+
+            current_face = faces[face_id]
+            face_landmarks = landmark_locator(image, current_face)
+            current_fl = search(face_landmarks)
+
+            forward_mask = np.ones_like(image).astype("uint8")
+            affine = compute_transformation_matrix(image, current_fl, False, target_face_scale=1.3)
+            aligned_face = warp(image, affine, output_shape=(512, 512, 3), preserve_range=True)
+            forward_mask = warp(
+                forward_mask, affine, output_shape=(512, 512, 3), order=0, preserve_range=True
+            )
+
+            affine_inverse = affine.inverse
+            cur_face = aligned_face
+            if replace_url != "":
+
+                face_name = x[:-4] + "_" + str(face_id + 1) + ".png"
+                cur_url = os.path.join(replace_url, face_name)
+                restored_face = Image.open(cur_url).convert("RGB")
+                restored_face = np.array(restored_face)
+                cur_face = restored_face
+
+            ## Histogram Color matching
+            A = cv2.cvtColor(aligned_face.astype("uint8"), cv2.COLOR_RGB2BGR)
+            B = cv2.cvtColor(cur_face.astype("uint8"), cv2.COLOR_RGB2BGR)
+            B = match_histograms(B, A)
+            cur_face = cv2.cvtColor(B.astype("uint8"), cv2.COLOR_BGR2RGB)
+
+            warped_back = warp(
+                cur_face,
+                affine_inverse,
+                output_shape=(origin_height, origin_width, 3),
+                order=3,
+                preserve_range=True,
+            )
+
+            backward_mask = warp(
+                forward_mask,
+                affine_inverse,
+                output_shape=(origin_height, origin_width, 3),
+                order=0,
+                preserve_range=True,
+            )  ## Nearest neighbour
+
+            blended = blur_blending_cv2(warped_back, blended, backward_mask)
+            blended *= 255.0
+
+        io.imsave(os.path.join(save_url, x), img_as_ubyte(blended / 255.0))
+
+        count += 1
+
+        if count % 1000 == 0:
+            print("%d have finished ..." % (count))
+

Face_Detection/detect_all_dlib.py

@@ -0,0 +1,184 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+import numpy as np
+import skimage.io as io
+
+# from FaceSDK.face_sdk import FaceDetection
+# from face_sdk import FaceDetection
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+from skimage.transform import SimilarityTransform
+from skimage.transform import warp
+from PIL import Image
+import torch.nn.functional as F
+import torchvision as tv
+import torchvision.utils as vutils
+import time
+import cv2
+import os
+from skimage import img_as_ubyte
+import json
+import argparse
+import dlib
+
+
+def _standard_face_pts():
+    pts = (
+        np.array([196.0, 226.0, 316.0, 226.0, 256.0, 286.0, 220.0, 360.4, 292.0, 360.4], np.float32) / 256.0
+        - 1.0
+    )
+
+    return np.reshape(pts, (5, 2))
+
+
+def _origin_face_pts():
+    pts = np.array([196.0, 226.0, 316.0, 226.0, 256.0, 286.0, 220.0, 360.4, 292.0, 360.4], np.float32)
+
+    return np.reshape(pts, (5, 2))
+
+
+def get_landmark(face_landmarks, id):
+    part = face_landmarks.part(id)
+    x = part.x
+    y = part.y
+
+    return (x, y)
+
+
+def search(face_landmarks):
+
+    x1, y1 = get_landmark(face_landmarks, 36)
+    x2, y2 = get_landmark(face_landmarks, 39)
+    x3, y3 = get_landmark(face_landmarks, 42)
+    x4, y4 = get_landmark(face_landmarks, 45)
+
+    x_nose, y_nose = get_landmark(face_landmarks, 30)
+
+    x_left_mouth, y_left_mouth = get_landmark(face_landmarks, 48)
+    x_right_mouth, y_right_mouth = get_landmark(face_landmarks, 54)
+
+    x_left_eye = int((x1 + x2) / 2)
+    y_left_eye = int((y1 + y2) / 2)
+    x_right_eye = int((x3 + x4) / 2)
+    y_right_eye = int((y3 + y4) / 2)
+
+    results = np.array(
+        [
+            [x_left_eye, y_left_eye],
+            [x_right_eye, y_right_eye],
+            [x_nose, y_nose],
+            [x_left_mouth, y_left_mouth],
+            [x_right_mouth, y_right_mouth],
+        ]
+    )
+
+    return results
+
+
+def compute_transformation_matrix(img, landmark, normalize, target_face_scale=1.0):
+
+    std_pts = _standard_face_pts()  # [-1,1]
+    target_pts = (std_pts * target_face_scale + 1) / 2 * 256.0
+
+    # print(target_pts)
+
+    h, w, c = img.shape
+    if normalize == True:
+        landmark[:, 0] = landmark[:, 0] / h * 2 - 1.0
+        landmark[:, 1] = landmark[:, 1] / w * 2 - 1.0
+
+    # print(landmark)
+
+    affine = SimilarityTransform()
+
+    affine.estimate(target_pts, landmark)
+
+    return affine.params
+
+
+def show_detection(image, box, landmark):
+    plt.imshow(image)
+    print(box[2] - box[0])
+    plt.gca().add_patch(
+        Rectangle(
+            (box[1], box[0]), box[2] - box[0], box[3] - box[1], linewidth=1, edgecolor="r", facecolor="none"
+        )
+    )
+    plt.scatter(landmark[0][0], landmark[0][1])
+    plt.scatter(landmark[1][0], landmark[1][1])
+    plt.scatter(landmark[2][0], landmark[2][1])
+    plt.scatter(landmark[3][0], landmark[3][1])
+    plt.scatter(landmark[4][0], landmark[4][1])
+    plt.show()
+
+
+def affine2theta(affine, input_w, input_h, target_w, target_h):
+    # param = np.linalg.inv(affine)
+    param = affine
+    theta = np.zeros([2, 3])
+    theta[0, 0] = param[0, 0] * input_h / target_h
+    theta[0, 1] = param[0, 1] * input_w / target_h
+    theta[0, 2] = (2 * param[0, 2] + param[0, 0] * input_h + param[0, 1] * input_w) / target_h - 1
+    theta[1, 0] = param[1, 0] * input_h / target_w
+    theta[1, 1] = param[1, 1] * input_w / target_w
+    theta[1, 2] = (2 * param[1, 2] + param[1, 0] * input_h + param[1, 1] * input_w) / target_w - 1
+    return theta
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--url", type=str, default="/home/jingliao/ziyuwan/celebrities", help="input")
+    parser.add_argument(
+        "--save_url", type=str, default="/home/jingliao/ziyuwan/celebrities_detected_face_reid", help="output"
+    )
+    opts = parser.parse_args()
+
+    url = opts.url
+    save_url = opts.save_url
+
+    ### If the origin url is None, then we don't need to reid the origin image
+
+    os.makedirs(url, exist_ok=True)
+    os.makedirs(save_url, exist_ok=True)
+
+    face_detector = dlib.get_frontal_face_detector()
+    landmark_locator = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
+
+    count = 0
+
+    map_id = {}
+    for x in os.listdir(url):
+        img_url = os.path.join(url, x)
+        pil_img = Image.open(img_url).convert("RGB")
+
+        image = np.array(pil_img)
+
+        start = time.time()
+        faces = face_detector(image)
+        done = time.time()
+
+        if len(faces) == 0:
+            print("Warning: There is no face in %s" % (x))
+            continue
+
+        print(len(faces))
+
+        if len(faces) > 0:
+            for face_id in range(len(faces)):
+                current_face = faces[face_id]
+                face_landmarks = landmark_locator(image, current_face)
+                current_fl = search(face_landmarks)
+
+                affine = compute_transformation_matrix(image, current_fl, False, target_face_scale=1.3)
+                aligned_face = warp(image, affine, output_shape=(256, 256, 3))
+                img_name = x[:-4] + "_" + str(face_id + 1)
+                io.imsave(os.path.join(save_url, img_name + ".png"), img_as_ubyte(aligned_face))
+
+        count += 1
+
+        if count % 1000 == 0:
+            print("%d have finished ..." % (count))
+

Face_Detection/detect_all_dlib_HR.py

@@ -0,0 +1,184 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+import numpy as np
+import skimage.io as io
+
+# from FaceSDK.face_sdk import FaceDetection
+# from face_sdk import FaceDetection
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+from skimage.transform import SimilarityTransform
+from skimage.transform import warp
+from PIL import Image
+import torch.nn.functional as F
+import torchvision as tv
+import torchvision.utils as vutils
+import time
+import cv2
+import os
+from skimage import img_as_ubyte
+import json
+import argparse
+import dlib
+
+
+def _standard_face_pts():
+    pts = (
+        np.array([196.0, 226.0, 316.0, 226.0, 256.0, 286.0, 220.0, 360.4, 292.0, 360.4], np.float32) / 256.0
+        - 1.0
+    )
+
+    return np.reshape(pts, (5, 2))
+
+
+def _origin_face_pts():
+    pts = np.array([196.0, 226.0, 316.0, 226.0, 256.0, 286.0, 220.0, 360.4, 292.0, 360.4], np.float32)
+
+    return np.reshape(pts, (5, 2))
+
+
+def get_landmark(face_landmarks, id):
+    part = face_landmarks.part(id)
+    x = part.x
+    y = part.y
+
+    return (x, y)
+
+
+def search(face_landmarks):
+
+    x1, y1 = get_landmark(face_landmarks, 36)
+    x2, y2 = get_landmark(face_landmarks, 39)
+    x3, y3 = get_landmark(face_landmarks, 42)
+    x4, y4 = get_landmark(face_landmarks, 45)
+
+    x_nose, y_nose = get_landmark(face_landmarks, 30)
+
+    x_left_mouth, y_left_mouth = get_landmark(face_landmarks, 48)
+    x_right_mouth, y_right_mouth = get_landmark(face_landmarks, 54)
+
+    x_left_eye = int((x1 + x2) / 2)
+    y_left_eye = int((y1 + y2) / 2)
+    x_right_eye = int((x3 + x4) / 2)
+    y_right_eye = int((y3 + y4) / 2)
+
+    results = np.array(
+        [
+            [x_left_eye, y_left_eye],
+            [x_right_eye, y_right_eye],
+            [x_nose, y_nose],
+            [x_left_mouth, y_left_mouth],
+            [x_right_mouth, y_right_mouth],
+        ]
+    )
+
+    return results
+
+
+def compute_transformation_matrix(img, landmark, normalize, target_face_scale=1.0):
+
+    std_pts = _standard_face_pts()  # [-1,1]
+    target_pts = (std_pts * target_face_scale + 1) / 2 * 512.0
+
+    # print(target_pts)
+
+    h, w, c = img.shape
+    if normalize == True:
+        landmark[:, 0] = landmark[:, 0] / h * 2 - 1.0
+        landmark[:, 1] = landmark[:, 1] / w * 2 - 1.0
+
+    # print(landmark)
+
+    affine = SimilarityTransform()
+
+    affine.estimate(target_pts, landmark)
+
+    return affine.params
+
+
+def show_detection(image, box, landmark):
+    plt.imshow(image)
+    print(box[2] - box[0])
+    plt.gca().add_patch(
+        Rectangle(
+            (box[1], box[0]), box[2] - box[0], box[3] - box[1], linewidth=1, edgecolor="r", facecolor="none"
+        )
+    )
+    plt.scatter(landmark[0][0], landmark[0][1])
+    plt.scatter(landmark[1][0], landmark[1][1])
+    plt.scatter(landmark[2][0], landmark[2][1])
+    plt.scatter(landmark[3][0], landmark[3][1])
+    plt.scatter(landmark[4][0], landmark[4][1])
+    plt.show()
+
+
+def affine2theta(affine, input_w, input_h, target_w, target_h):
+    # param = np.linalg.inv(affine)
+    param = affine
+    theta = np.zeros([2, 3])
+    theta[0, 0] = param[0, 0] * input_h / target_h
+    theta[0, 1] = param[0, 1] * input_w / target_h
+    theta[0, 2] = (2 * param[0, 2] + param[0, 0] * input_h + param[0, 1] * input_w) / target_h - 1
+    theta[1, 0] = param[1, 0] * input_h / target_w
+    theta[1, 1] = param[1, 1] * input_w / target_w
+    theta[1, 2] = (2 * param[1, 2] + param[1, 0] * input_h + param[1, 1] * input_w) / target_w - 1
+    return theta
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--url", type=str, default="/home/jingliao/ziyuwan/celebrities", help="input")
+    parser.add_argument(
+        "--save_url", type=str, default="/home/jingliao/ziyuwan/celebrities_detected_face_reid", help="output"
+    )
+    opts = parser.parse_args()
+
+    url = opts.url
+    save_url = opts.save_url
+
+    ### If the origin url is None, then we don't need to reid the origin image
+
+    os.makedirs(url, exist_ok=True)
+    os.makedirs(save_url, exist_ok=True)
+
+    face_detector = dlib.get_frontal_face_detector()
+    landmark_locator = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
+
+    count = 0
+
+    map_id = {}
+    for x in os.listdir(url):
+        img_url = os.path.join(url, x)
+        pil_img = Image.open(img_url).convert("RGB")
+
+        image = np.array(pil_img)
+
+        start = time.time()
+        faces = face_detector(image)
+        done = time.time()
+
+        if len(faces) == 0:
+            print("Warning: There is no face in %s" % (x))
+            continue
+
+        print(len(faces))
+
+        if len(faces) > 0:
+            for face_id in range(len(faces)):
+                current_face = faces[face_id]
+                face_landmarks = landmark_locator(image, current_face)
+                current_fl = search(face_landmarks)
+
+                affine = compute_transformation_matrix(image, current_fl, False, target_face_scale=1.3)
+                aligned_face = warp(image, affine, output_shape=(512, 512, 3))
+                img_name = x[:-4] + "_" + str(face_id + 1)
+                io.imsave(os.path.join(save_url, img_name + ".png"), img_as_ubyte(aligned_face))
+
+        count += 1
+
+        if count % 1000 == 0:
+            print("%d have finished ..." % (count))
+

Face_Enhancement/data/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import importlib
+import torch.utils.data
+from data.base_dataset import BaseDataset
+from data.face_dataset import FaceTestDataset
+
+
+def create_dataloader(opt):
+
+    instance = FaceTestDataset()
+    instance.initialize(opt)
+    print("dataset [%s] of size %d was created" % (type(instance).__name__, len(instance)))
+    dataloader = torch.utils.data.DataLoader(
+        instance,
+        batch_size=opt.batchSize,
+        shuffle=not opt.serial_batches,
+        num_workers=int(opt.nThreads),
+        drop_last=opt.isTrain,
+    )
+    return dataloader

Face_Enhancement/data/base_dataset.py

@@ -0,0 +1,125 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch.utils.data as data
+from PIL import Image
+import torchvision.transforms as transforms
+import numpy as np
+import random
+
+
+class BaseDataset(data.Dataset):
+    def __init__(self):
+        super(BaseDataset, self).__init__()
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        return parser
+
+    def initialize(self, opt):
+        pass
+
+
+def get_params(opt, size):
+    w, h = size
+    new_h = h
+    new_w = w
+    if opt.preprocess_mode == "resize_and_crop":
+        new_h = new_w = opt.load_size
+    elif opt.preprocess_mode == "scale_width_and_crop":
+        new_w = opt.load_size
+        new_h = opt.load_size * h // w
+    elif opt.preprocess_mode == "scale_shortside_and_crop":
+        ss, ls = min(w, h), max(w, h)  # shortside and longside
+        width_is_shorter = w == ss
+        ls = int(opt.load_size * ls / ss)
+        new_w, new_h = (ss, ls) if width_is_shorter else (ls, ss)
+
+    x = random.randint(0, np.maximum(0, new_w - opt.crop_size))
+    y = random.randint(0, np.maximum(0, new_h - opt.crop_size))
+
+    flip = random.random() > 0.5
+    return {"crop_pos": (x, y), "flip": flip}
+
+
+def get_transform(opt, params, method=Image.BICUBIC, normalize=True, toTensor=True):
+    transform_list = []
+    if "resize" in opt.preprocess_mode:
+        osize = [opt.load_size, opt.load_size]
+        transform_list.append(transforms.Resize(osize, interpolation=method))
+    elif "scale_width" in opt.preprocess_mode:
+        transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.load_size, method)))
+    elif "scale_shortside" in opt.preprocess_mode:
+        transform_list.append(transforms.Lambda(lambda img: __scale_shortside(img, opt.load_size, method)))
+
+    if "crop" in opt.preprocess_mode:
+        transform_list.append(transforms.Lambda(lambda img: __crop(img, params["crop_pos"], opt.crop_size)))
+
+    if opt.preprocess_mode == "none":
+        base = 32
+        transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method)))
+
+    if opt.preprocess_mode == "fixed":
+        w = opt.crop_size
+        h = round(opt.crop_size / opt.aspect_ratio)
+        transform_list.append(transforms.Lambda(lambda img: __resize(img, w, h, method)))
+
+    if opt.isTrain and not opt.no_flip:
+        transform_list.append(transforms.Lambda(lambda img: __flip(img, params["flip"])))
+
+    if toTensor:
+        transform_list += [transforms.ToTensor()]
+
+    if normalize:
+        transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
+    return transforms.Compose(transform_list)
+
+
+def normalize():
+    return transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+
+
+def __resize(img, w, h, method=Image.BICUBIC):
+    return img.resize((w, h), method)
+
+
+def __make_power_2(img, base, method=Image.BICUBIC):
+    ow, oh = img.size
+    h = int(round(oh / base) * base)
+    w = int(round(ow / base) * base)
+    if (h == oh) and (w == ow):
+        return img
+    return img.resize((w, h), method)
+
+
+def __scale_width(img, target_width, method=Image.BICUBIC):
+    ow, oh = img.size
+    if ow == target_width:
+        return img
+    w = target_width
+    h = int(target_width * oh / ow)
+    return img.resize((w, h), method)
+
+
+def __scale_shortside(img, target_width, method=Image.BICUBIC):
+    ow, oh = img.size
+    ss, ls = min(ow, oh), max(ow, oh)  # shortside and longside
+    width_is_shorter = ow == ss
+    if ss == target_width:
+        return img
+    ls = int(target_width * ls / ss)
+    nw, nh = (ss, ls) if width_is_shorter else (ls, ss)
+    return img.resize((nw, nh), method)
+
+
+def __crop(img, pos, size):
+    ow, oh = img.size
+    x1, y1 = pos
+    tw = th = size
+    return img.crop((x1, y1, x1 + tw, y1 + th))
+
+
+def __flip(img, flip):
+    if flip:
+        return img.transpose(Image.FLIP_LEFT_RIGHT)
+    return img

Face_Enhancement/data/custom_dataset.py

@@ -0,0 +1,56 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from data.pix2pix_dataset import Pix2pixDataset
+from data.image_folder import make_dataset
+
+
+class CustomDataset(Pix2pixDataset):
+    """ Dataset that loads images from directories
+        Use option --label_dir, --image_dir, --instance_dir to specify the directories.
+        The images in the directories are sorted in alphabetical order and paired in order.
+    """
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        parser = Pix2pixDataset.modify_commandline_options(parser, is_train)
+        parser.set_defaults(preprocess_mode="resize_and_crop")
+        load_size = 286 if is_train else 256
+        parser.set_defaults(load_size=load_size)
+        parser.set_defaults(crop_size=256)
+        parser.set_defaults(display_winsize=256)
+        parser.set_defaults(label_nc=13)
+        parser.set_defaults(contain_dontcare_label=False)
+
+        parser.add_argument(
+            "--label_dir", type=str, required=True, help="path to the directory that contains label images"
+        )
+        parser.add_argument(
+            "--image_dir", type=str, required=True, help="path to the directory that contains photo images"
+        )
+        parser.add_argument(
+            "--instance_dir",
+            type=str,
+            default="",
+            help="path to the directory that contains instance maps. Leave black if not exists",
+        )
+        return parser
+
+    def get_paths(self, opt):
+        label_dir = opt.label_dir
+        label_paths = make_dataset(label_dir, recursive=False, read_cache=True)
+
+        image_dir = opt.image_dir
+        image_paths = make_dataset(image_dir, recursive=False, read_cache=True)
+
+        if len(opt.instance_dir) > 0:
+            instance_dir = opt.instance_dir
+            instance_paths = make_dataset(instance_dir, recursive=False, read_cache=True)
+        else:
+            instance_paths = []
+
+        assert len(label_paths) == len(
+            image_paths
+        ), "The #images in %s and %s do not match. Is there something wrong?"
+
+        return label_paths, image_paths, instance_paths

Face_Enhancement/data/face_dataset.py

@@ -0,0 +1,102 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from data.base_dataset import BaseDataset, get_params, get_transform
+from PIL import Image
+import util.util as util
+import os
+import torch
+
+
+class FaceTestDataset(BaseDataset):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        parser.add_argument(
+            "--no_pairing_check",
+            action="store_true",
+            help="If specified, skip sanity check of correct label-image file pairing",
+        )
+        #    parser.set_defaults(contain_dontcare_label=False)
+        #    parser.set_defaults(no_instance=True)
+        return parser
+
+    def initialize(self, opt):
+        self.opt = opt
+
+        image_path = os.path.join(opt.dataroot, opt.old_face_folder)
+        label_path = os.path.join(opt.dataroot, opt.old_face_label_folder)
+
+        image_list = os.listdir(image_path)
+        image_list = sorted(image_list)
+        # image_list=image_list[:opt.max_dataset_size]
+
+        self.label_paths = label_path  ## Just the root dir
+        self.image_paths = image_list  ## All the image name
+
+        self.parts = [
+            "skin",
+            "hair",
+            "l_brow",
+            "r_brow",
+            "l_eye",
+            "r_eye",
+            "eye_g",
+            "l_ear",
+            "r_ear",
+            "ear_r",
+            "nose",
+            "mouth",
+            "u_lip",
+            "l_lip",
+            "neck",
+            "neck_l",
+            "cloth",
+            "hat",
+        ]
+
+        size = len(self.image_paths)
+        self.dataset_size = size
+
+    def __getitem__(self, index):
+
+        params = get_params(self.opt, (-1, -1))
+        image_name = self.image_paths[index]
+        image_path = os.path.join(self.opt.dataroot, self.opt.old_face_folder, image_name)
+        image = Image.open(image_path)
+        image = image.convert("RGB")
+
+        transform_image = get_transform(self.opt, params)
+        image_tensor = transform_image(image)
+
+        img_name = image_name[:-4]
+        transform_label = get_transform(self.opt, params, method=Image.NEAREST, normalize=False)
+        full_label = []
+
+        cnt = 0
+
+        for each_part in self.parts:
+            part_name = img_name + "_" + each_part + ".png"
+            part_url = os.path.join(self.label_paths, part_name)
+
+            if os.path.exists(part_url):
+                label = Image.open(part_url).convert("RGB")
+                label_tensor = transform_label(label)  ## 3 channels and pixel [0,1]
+                full_label.append(label_tensor[0])
+            else:
+                current_part = torch.zeros((self.opt.load_size, self.opt.load_size))
+                full_label.append(current_part)
+                cnt += 1
+
+        full_label_tensor = torch.stack(full_label, 0)
+
+        input_dict = {
+            "label": full_label_tensor,
+            "image": image_tensor,
+            "path": image_path,
+        }
+
+        return input_dict
+
+    def __len__(self):
+        return self.dataset_size
+

Face_Enhancement/data/image_folder.py

@@ -0,0 +1,101 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch.utils.data as data
+from PIL import Image
+import os
+
+IMG_EXTENSIONS = [
+    ".jpg",
+    ".JPG",
+    ".jpeg",
+    ".JPEG",
+    ".png",
+    ".PNG",
+    ".ppm",
+    ".PPM",
+    ".bmp",
+    ".BMP",
+    ".tiff",
+    ".webp",
+]
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset_rec(dir, images):
+    assert os.path.isdir(dir), "%s is not a valid directory" % dir
+
+    for root, dnames, fnames in sorted(os.walk(dir, followlinks=True)):
+        for fname in fnames:
+            if is_image_file(fname):
+                path = os.path.join(root, fname)
+                images.append(path)
+
+
+def make_dataset(dir, recursive=False, read_cache=False, write_cache=False):
+    images = []
+
+    if read_cache:
+        possible_filelist = os.path.join(dir, "files.list")
+        if os.path.isfile(possible_filelist):
+            with open(possible_filelist, "r") as f:
+                images = f.read().splitlines()
+                return images
+
+    if recursive:
+        make_dataset_rec(dir, images)
+    else:
+        assert os.path.isdir(dir) or os.path.islink(dir), "%s is not a valid directory" % dir
+
+        for root, dnames, fnames in sorted(os.walk(dir)):
+            for fname in fnames:
+                if is_image_file(fname):
+                    path = os.path.join(root, fname)
+                    images.append(path)
+
+    if write_cache:
+        filelist_cache = os.path.join(dir, "files.list")
+        with open(filelist_cache, "w") as f:
+            for path in images:
+                f.write("%s\n" % path)
+            print("wrote filelist cache at %s" % filelist_cache)
+
+    return images
+
+
+def default_loader(path):
+    return Image.open(path).convert("RGB")
+
+
+class ImageFolder(data.Dataset):
+    def __init__(self, root, transform=None, return_paths=False, loader=default_loader):
+        imgs = make_dataset(root)
+        if len(imgs) == 0:
+            raise (
+                RuntimeError(
+                    "Found 0 images in: " + root + "\n"
+                    "Supported image extensions are: " + ",".join(IMG_EXTENSIONS)
+                )
+            )
+
+        self.root = root
+        self.imgs = imgs
+        self.transform = transform
+        self.return_paths = return_paths
+        self.loader = loader
+
+    def __getitem__(self, index):
+        path = self.imgs[index]
+        img = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.return_paths:
+            return img, path
+        else:
+            return img
+
+    def __len__(self):
+        return len(self.imgs)

Face_Enhancement/data/pix2pix_dataset.py

@@ -0,0 +1,108 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from data.base_dataset import BaseDataset, get_params, get_transform
+from PIL import Image
+import util.util as util
+import os
+
+
+class Pix2pixDataset(BaseDataset):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        parser.add_argument(
+            "--no_pairing_check",
+            action="store_true",
+            help="If specified, skip sanity check of correct label-image file pairing",
+        )
+        return parser
+
+    def initialize(self, opt):
+        self.opt = opt
+
+        label_paths, image_paths, instance_paths = self.get_paths(opt)
+
+        util.natural_sort(label_paths)
+        util.natural_sort(image_paths)
+        if not opt.no_instance:
+            util.natural_sort(instance_paths)
+
+        label_paths = label_paths[: opt.max_dataset_size]
+        image_paths = image_paths[: opt.max_dataset_size]
+        instance_paths = instance_paths[: opt.max_dataset_size]
+
+        if not opt.no_pairing_check:
+            for path1, path2 in zip(label_paths, image_paths):
+                assert self.paths_match(path1, path2), (
+                    "The label-image pair (%s, %s) do not look like the right pair because the filenames are quite different. Are you sure about the pairing? Please see data/pix2pix_dataset.py to see what is going on, and use --no_pairing_check to bypass this."
+                    % (path1, path2)
+                )
+
+        self.label_paths = label_paths
+        self.image_paths = image_paths
+        self.instance_paths = instance_paths
+
+        size = len(self.label_paths)
+        self.dataset_size = size
+
+    def get_paths(self, opt):
+        label_paths = []
+        image_paths = []
+        instance_paths = []
+        assert False, "A subclass of Pix2pixDataset must override self.get_paths(self, opt)"
+        return label_paths, image_paths, instance_paths
+
+    def paths_match(self, path1, path2):
+        filename1_without_ext = os.path.splitext(os.path.basename(path1))[0]
+        filename2_without_ext = os.path.splitext(os.path.basename(path2))[0]
+        return filename1_without_ext == filename2_without_ext
+
+    def __getitem__(self, index):
+        # Label Image
+        label_path = self.label_paths[index]
+        label = Image.open(label_path)
+        params = get_params(self.opt, label.size)
+        transform_label = get_transform(self.opt, params, method=Image.NEAREST, normalize=False)
+        label_tensor = transform_label(label) * 255.0
+        label_tensor[label_tensor == 255] = self.opt.label_nc  # 'unknown' is opt.label_nc
+
+        # input image (real images)
+        image_path = self.image_paths[index]
+        assert self.paths_match(
+            label_path, image_path
+        ), "The label_path %s and image_path %s don't match." % (label_path, image_path)
+        image = Image.open(image_path)
+        image = image.convert("RGB")
+
+        transform_image = get_transform(self.opt, params)
+        image_tensor = transform_image(image)
+
+        # if using instance maps
+        if self.opt.no_instance:
+            instance_tensor = 0
+        else:
+            instance_path = self.instance_paths[index]
+            instance = Image.open(instance_path)
+            if instance.mode == "L":
+                instance_tensor = transform_label(instance) * 255
+                instance_tensor = instance_tensor.long()
+            else:
+                instance_tensor = transform_label(instance)
+
+        input_dict = {
+            "label": label_tensor,
+            "instance": instance_tensor,
+            "image": image_tensor,
+            "path": image_path,
+        }
+
+        # Give subclasses a chance to modify the final output
+        self.postprocess(input_dict)
+
+        return input_dict
+
+    def postprocess(self, input_dict):
+        return input_dict
+
+    def __len__(self):
+        return self.dataset_size

Face_Enhancement/models/__init__.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import importlib
+import torch
+
+
+def find_model_using_name(model_name):
+    # Given the option --model [modelname],
+    # the file "models/modelname_model.py"
+    # will be imported.
+    model_filename = "models." + model_name + "_model"
+    modellib = importlib.import_module(model_filename)
+
+    # In the file, the class called ModelNameModel() will
+    # be instantiated. It has to be a subclass of torch.nn.Module,
+    # and it is case-insensitive.
+    model = None
+    target_model_name = model_name.replace("_", "") + "model"
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower() and issubclass(cls, torch.nn.Module):
+            model = cls
+
+    if model is None:
+        print(
+            "In %s.py, there should be a subclass of torch.nn.Module with class name that matches %s in lowercase."
+            % (model_filename, target_model_name)
+        )
+        exit(0)
+
+    return model
+
+
+def get_option_setter(model_name):
+    model_class = find_model_using_name(model_name)
+    return model_class.modify_commandline_options
+
+
+def create_model(opt):
+    model = find_model_using_name(opt.model)
+    instance = model(opt)
+    print("model [%s] was created" % (type(instance).__name__))
+
+    return instance

Face_Enhancement/models/networks/Synchronized-BatchNorm-PyTorch/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2018 Jiayuan MAO
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Face_Enhancement/models/networks/Synchronized-BatchNorm-PyTorch/README.md

@@ -0,0 +1,118 @@
+# Synchronized-BatchNorm-PyTorch
+
+**IMPORTANT: Please read the "Implementation details and highlights" section before use.**
+
+Synchronized Batch Normalization implementation in PyTorch.
+
+This module differs from the built-in PyTorch BatchNorm as the mean and
+standard-deviation are reduced across all devices during training.
+
+For example, when one uses `nn.DataParallel` to wrap the network during
+training, PyTorch's implementation normalize the tensor on each device using
+the statistics only on that device, which accelerated the computation and
+is also easy to implement, but the statistics might be inaccurate.
+Instead, in this synchronized version, the statistics will be computed
+over all training samples distributed on multiple devices.
+
+Note that, for one-GPU or CPU-only case, this module behaves exactly same
+as the built-in PyTorch implementation.
+
+This module is currently only a prototype version for research usages. As mentioned below,
+it has its limitations and may even suffer from some design problems. If you have any
+questions or suggestions, please feel free to
+[open an issue](https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues) or 
+[submit a pull request](https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues).
+
+## Why Synchronized BatchNorm?
+
+Although the typical implementation of BatchNorm working on multiple devices (GPUs)
+is fast (with no communication overhead), it inevitably reduces the size of batch size,
+which potentially degenerates the performance. This is not a significant issue in some
+standard vision tasks such as ImageNet classification (as the batch size per device
+is usually large enough to obtain good statistics). However, it will hurt the performance
+in some tasks that the batch size is usually very small (e.g., 1 per GPU).
+
+For example, the importance of synchronized batch normalization in object detection has been recently proved with a
+an extensive analysis in the paper [MegDet: A Large Mini-Batch Object Detector](https://arxiv.org/abs/1711.07240).
+
+## Usage
+
+To use the Synchronized Batch Normalization, we add a data parallel replication callback. This introduces a slight
+difference with typical usage of the `nn.DataParallel`.
+
+Use it with a provided, customized data parallel wrapper:
+
+```python
+from sync_batchnorm import SynchronizedBatchNorm1d, DataParallelWithCallback
+
+sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+```
+
+Or, if you are using a customized data parallel module, you can use this library as a monkey patching.
+
+```python
+from torch.nn import DataParallel  # or your customized DataParallel module
+from sync_batchnorm import SynchronizedBatchNorm1d, patch_replication_callback
+
+sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+sync_bn = DataParallel(sync_bn, device_ids=[0, 1])
+patch_replication_callback(sync_bn)  # monkey-patching
+```
+
+You can use `convert_model` to convert your model to use Synchronized BatchNorm easily.
+
+```python
+import torch.nn as nn
+from torchvision import models
+from sync_batchnorm import convert_model
+# m is a standard pytorch model
+m = models.resnet18(True)
+m = nn.DataParallel(m)
+# after convert, m is using SyncBN
+m = convert_model(m)
+```
+
+See also `tests/test_sync_batchnorm.py` for numeric result comparison.
+
+## Implementation details and highlights
+
+If you are interested in how batch statistics are reduced and broadcasted among multiple devices, please take a look
+at the code with detailed comments. Here we only emphasize some highlights of the implementation:
+
+- This implementation is in pure-python. No C++ extra extension libs.
+- Easy to use as demonstrated above.
+- It uses unbiased variance to update the moving average, and use `sqrt(max(var, eps))` instead of `sqrt(var + eps)`.
+- The implementation requires that each module on different devices should invoke the `batchnorm` for exactly SAME
+amount of times in each forward pass. For example, you can not only call `batchnorm` on GPU0 but not on GPU1. The `#i
+(i = 1, 2, 3, ...)` calls of the `batchnorm` on each device will be viewed as a whole and the statistics will be reduced.
+This is tricky but is a good way to handle PyTorch's dynamic computation graph. Although sounds complicated, this
+will usually not be the issue for most of the models.
+
+## Known issues
+
+#### Runtime error on backward pass.
+
+Due to a [PyTorch Bug](https://github.com/pytorch/pytorch/issues/3883), using old PyTorch libraries will trigger an `RuntimeError` with messages like:
+
+```
+Assertion `pos >= 0 && pos < buffer.size()` failed.
+```
+
+This has already been solved in the newest PyTorch repo, which, unfortunately, has not been pushed to the official and anaconda binary release. Thus, you are required to build the PyTorch package from the source according to the
+ instructions [here](https://github.com/pytorch/pytorch#from-source).
+
+#### Numeric error.
+
+Because this library does not fuse the normalization and statistics operations in C++ (nor CUDA), it is less
+numerically stable compared to the original PyTorch implementation. Detailed analysis can be found in
+`tests/test_sync_batchnorm.py`.
+
+## Authors and License:
+
+Copyright (c) 2018-, [Jiayuan Mao](https://vccy.xyz).
+
+**Contributors**: [Tete Xiao](https://tetexiao.com), [DTennant](https://github.com/DTennant).
+
+Distributed under **MIT License** (See LICENSE)
+

Face_Enhancement/models/networks/Synchronized-BatchNorm-PyTorch/sync_batchnorm/__init__.py

@@ -0,0 +1,14 @@
+# -*- coding: utf-8 -*-
+# File   : __init__.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+from .batchnorm import set_sbn_eps_mode
+from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d
+from .batchnorm import patch_sync_batchnorm, convert_model
+from .replicate import DataParallelWithCallback, patch_replication_callback

Face_Enhancement/models/networks/Synchronized-BatchNorm-PyTorch/sync_batchnorm/batchnorm.py

@@ -0,0 +1,412 @@
+# -*- coding: utf-8 -*-
+# File   : batchnorm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import collections
+import contextlib
+
+import torch
+import torch.nn.functional as F
+
+from torch.nn.modules.batchnorm import _BatchNorm
+
+try:
+    from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast
+except ImportError:
+    ReduceAddCoalesced = Broadcast = None
+
+try:
+    from jactorch.parallel.comm import SyncMaster
+    from jactorch.parallel.data_parallel import JacDataParallel as DataParallelWithCallback
+except ImportError:
+    from .comm import SyncMaster
+    from .replicate import DataParallelWithCallback
+
+__all__ = [
+    'set_sbn_eps_mode',
+    'SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d', 'SynchronizedBatchNorm3d',
+    'patch_sync_batchnorm', 'convert_model'
+]
+
+
+SBN_EPS_MODE = 'clamp'
+
+
+def set_sbn_eps_mode(mode):
+    global SBN_EPS_MODE
+    assert mode in ('clamp', 'plus')
+    SBN_EPS_MODE = mode
+
+
+def _sum_ft(tensor):
+    """sum over the first and last dimention"""
+    return tensor.sum(dim=0).sum(dim=-1)
+
+
+def _unsqueeze_ft(tensor):
+    """add new dimensions at the front and the tail"""
+    return tensor.unsqueeze(0).unsqueeze(-1)
+
+
+_ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size'])
+_MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std'])
+
+
+class _SynchronizedBatchNorm(_BatchNorm):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True):
+        assert ReduceAddCoalesced is not None, 'Can not use Synchronized Batch Normalization without CUDA support.'
+
+        super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine,
+                                                     track_running_stats=track_running_stats)
+
+        if not self.track_running_stats:
+            import warnings
+            warnings.warn('track_running_stats=False is not supported by the SynchronizedBatchNorm.')
+
+        self._sync_master = SyncMaster(self._data_parallel_master)
+
+        self._is_parallel = False
+        self._parallel_id = None
+        self._slave_pipe = None
+
+    def forward(self, input):
+        # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation.
+        if not (self._is_parallel and self.training):
+            return F.batch_norm(
+                input, self.running_mean, self.running_var, self.weight, self.bias,
+                self.training, self.momentum, self.eps)
+
+        # Resize the input to (B, C, -1).
+        input_shape = input.size()
+        assert input.size(1) == self.num_features, 'Channel size mismatch: got {}, expect {}.'.format(input.size(1), self.num_features)
+        input = input.view(input.size(0), self.num_features, -1)
+
+        # Compute the sum and square-sum.
+        sum_size = input.size(0) * input.size(2)
+        input_sum = _sum_ft(input)
+        input_ssum = _sum_ft(input ** 2)
+
+        # Reduce-and-broadcast the statistics.
+        if self._parallel_id == 0:
+            mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size))
+        else:
+            mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size))
+
+        # Compute the output.
+        if self.affine:
+            # MJY:: Fuse the multiplication for speed.
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias)
+        else:
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std)
+
+        # Reshape it.
+        return output.view(input_shape)
+
+    def __data_parallel_replicate__(self, ctx, copy_id):
+        self._is_parallel = True
+        self._parallel_id = copy_id
+
+        # parallel_id == 0 means master device.
+        if self._parallel_id == 0:
+            ctx.sync_master = self._sync_master
+        else:
+            self._slave_pipe = ctx.sync_master.register_slave(copy_id)
+
+    def _data_parallel_master(self, intermediates):
+        """Reduce the sum and square-sum, compute the statistics, and broadcast it."""
+
+        # Always using same "device order" makes the ReduceAdd operation faster.
+        # Thanks to:: Tete Xiao (http://tetexiao.com/)
+        intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device())
+
+        to_reduce = [i[1][:2] for i in intermediates]
+        to_reduce = [j for i in to_reduce for j in i]  # flatten
+        target_gpus = [i[1].sum.get_device() for i in intermediates]
+
+        sum_size = sum([i[1].sum_size for i in intermediates])
+        sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce)
+        mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size)
+
+        broadcasted = Broadcast.apply(target_gpus, mean, inv_std)
+
+        outputs = []
+        for i, rec in enumerate(intermediates):
+            outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2])))
+
+        return outputs
+
+    def _compute_mean_std(self, sum_, ssum, size):
+        """Compute the mean and standard-deviation with sum and square-sum. This method
+        also maintains the moving average on the master device."""
+        assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.'
+        mean = sum_ / size
+        sumvar = ssum - sum_ * mean
+        unbias_var = sumvar / (size - 1)
+        bias_var = sumvar / size
+
+        if hasattr(torch, 'no_grad'):
+            with torch.no_grad():
+                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data
+                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data
+        else:
+            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data
+            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data
+
+        if SBN_EPS_MODE == 'clamp':
+            return mean, bias_var.clamp(self.eps) ** -0.5
+        elif SBN_EPS_MODE == 'plus':
+            return mean, (bias_var + self.eps) ** -0.5
+        else:
+            raise ValueError('Unknown EPS mode: {}.'.format(SBN_EPS_MODE))
+
+
+class SynchronizedBatchNorm1d(_SynchronizedBatchNorm):
+    r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a
+    mini-batch.
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm1d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of size
+            `batch_size x num_features [x width]`
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape::
+        - Input: :math:`(N, C)` or :math:`(N, C, L)`
+        - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 2 and input.dim() != 3:
+            raise ValueError('expected 2D or 3D input (got {}D input)'
+                             .format(input.dim()))
+
+
+class SynchronizedBatchNorm2d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch
+    of 3d inputs
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm2d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape::
+        - Input: :math:`(N, C, H, W)`
+        - Output: :math:`(N, C, H, W)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 4:
+            raise ValueError('expected 4D input (got {}D input)'
+                             .format(input.dim()))
+
+
+class SynchronizedBatchNorm3d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch
+    of 4d inputs
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm3d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm
+    or Spatio-temporal BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x depth x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape::
+        - Input: :math:`(N, C, D, H, W)`
+        - Output: :math:`(N, C, D, H, W)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 5:
+            raise ValueError('expected 5D input (got {}D input)'
+                             .format(input.dim()))
+
+
+@contextlib.contextmanager
+def patch_sync_batchnorm():
+    import torch.nn as nn
+
+    backup = nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d
+
+    nn.BatchNorm1d = SynchronizedBatchNorm1d
+    nn.BatchNorm2d = SynchronizedBatchNorm2d
+    nn.BatchNorm3d = SynchronizedBatchNorm3d
+
+    yield
+
+    nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d = backup
+
+
+def convert_model(module):
+    """Traverse the input module and its child recursively
+       and replace all instance of torch.nn.modules.batchnorm.BatchNorm*N*d
+       to SynchronizedBatchNorm*N*d
+
+    Args:
+        module: the input module needs to be convert to SyncBN model
+
+    Examples:
+        >>> import torch.nn as nn
+        >>> import torchvision
+        >>> # m is a standard pytorch model
+        >>> m = torchvision.models.resnet18(True)
+        >>> m = nn.DataParallel(m)
+        >>> # after convert, m is using SyncBN
+        >>> m = convert_model(m)
+    """
+    if isinstance(module, torch.nn.DataParallel):
+        mod = module.module
+        mod = convert_model(mod)
+        mod = DataParallelWithCallback(mod, device_ids=module.device_ids)
+        return mod
+
+    mod = module
+    for pth_module, sync_module in zip([torch.nn.modules.batchnorm.BatchNorm1d,
+                                        torch.nn.modules.batchnorm.BatchNorm2d,
+                                        torch.nn.modules.batchnorm.BatchNorm3d],
+                                       [SynchronizedBatchNorm1d,
+                                        SynchronizedBatchNorm2d,
+                                        SynchronizedBatchNorm3d]):
+        if isinstance(module, pth_module):
+            mod = sync_module(module.num_features, module.eps, module.momentum, module.affine)
+            mod.running_mean = module.running_mean
+            mod.running_var = module.running_var
+            if module.affine:
+                mod.weight.data = module.weight.data.clone().detach()
+                mod.bias.data = module.bias.data.clone().detach()
+
+    for name, child in module.named_children():
+        mod.add_module(name, convert_model(child))
+
+    return mod

Face_Enhancement/models/networks/Synchronized-BatchNorm-PyTorch/sync_batchnorm/batchnorm_reimpl.py

@@ -0,0 +1,74 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : batchnorm_reimpl.py
+# Author : acgtyrant
+# Date   : 11/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+
+__all__ = ['BatchNorm2dReimpl']
+
+
+class BatchNorm2dReimpl(nn.Module):
+    """
+    A re-implementation of batch normalization, used for testing the numerical
+    stability.
+
+    Author: acgtyrant
+    See also:
+    https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/14
+    """
+    def __init__(self, num_features, eps=1e-5, momentum=0.1):
+        super().__init__()
+
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.weight = nn.Parameter(torch.empty(num_features))
+        self.bias = nn.Parameter(torch.empty(num_features))
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+        self.reset_parameters()
+
+    def reset_running_stats(self):
+        self.running_mean.zero_()
+        self.running_var.fill_(1)
+
+    def reset_parameters(self):
+        self.reset_running_stats()
+        init.uniform_(self.weight)
+        init.zeros_(self.bias)
+
+    def forward(self, input_):
+        batchsize, channels, height, width = input_.size()
+        numel = batchsize * height * width
+        input_ = input_.permute(1, 0, 2, 3).contiguous().view(channels, numel)
+        sum_ = input_.sum(1)
+        sum_of_square = input_.pow(2).sum(1)
+        mean = sum_ / numel
+        sumvar = sum_of_square - sum_ * mean
+
+        self.running_mean = (
+                (1 - self.momentum) * self.running_mean
+                + self.momentum * mean.detach()
+        )
+        unbias_var = sumvar / (numel - 1)
+        self.running_var = (
+                (1 - self.momentum) * self.running_var
+                + self.momentum * unbias_var.detach()
+        )
+
+        bias_var = sumvar / numel
+        inv_std = 1 / (bias_var + self.eps).pow(0.5)
+        output = (
+                (input_ - mean.unsqueeze(1)) * inv_std.unsqueeze(1) *
+                self.weight.unsqueeze(1) + self.bias.unsqueeze(1))
+
+        return output.view(channels, batchsize, height, width).permute(1, 0, 2, 3).contiguous()
+

Face_Enhancement/models/networks/Synchronized-BatchNorm-PyTorch/sync_batchnorm/comm.py

@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+# File   : comm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import queue
+import collections
+import threading
+
+__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster']
+
+
+class FutureResult(object):
+    """A thread-safe future implementation. Used only as one-to-one pipe."""
+
+    def __init__(self):
+        self._result = None
+        self._lock = threading.Lock()
+        self._cond = threading.Condition(self._lock)
+
+    def put(self, result):
+        with self._lock:
+            assert self._result is None, 'Previous result has\'t been fetched.'
+            self._result = result
+            self._cond.notify()
+
+    def get(self):
+        with self._lock:
+            if self._result is None:
+                self._cond.wait()
+
+            res = self._result
+            self._result = None
+            return res
+
+
+_MasterRegistry = collections.namedtuple('MasterRegistry', ['result'])
+_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result'])
+
+
+class SlavePipe(_SlavePipeBase):
+    """Pipe for master-slave communication."""
+
+    def run_slave(self, msg):
+        self.queue.put((self.identifier, msg))
+        ret = self.result.get()
+        self.queue.put(True)
+        return ret
+
+
+class SyncMaster(object):
+    """An abstract `SyncMaster` object.
+
+    - During the replication, as the data parallel will trigger an callback of each module, all slave devices should
+    call `register(id)` and obtain an `SlavePipe` to communicate with the master.
+    - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected,
+    and passed to a registered callback.
+    - After receiving the messages, the master device should gather the information and determine to message passed
+    back to each slave devices.
+    """
+
+    def __init__(self, master_callback):
+        """
+
+        Args:
+            master_callback: a callback to be invoked after having collected messages from slave devices.
+        """
+        self._master_callback = master_callback
+        self._queue = queue.Queue()
+        self._registry = collections.OrderedDict()
+        self._activated = False
+
+    def __getstate__(self):
+        return {'master_callback': self._master_callback}
+
+    def __setstate__(self, state):
+        self.__init__(state['master_callback'])
+
+    def register_slave(self, identifier):
+        """
+        Register an slave device.
+
+        Args:
+            identifier: an identifier, usually is the device id.
+
+        Returns: a `SlavePipe` object which can be used to communicate with the master device.
+
+        """
+        if self._activated:
+            assert self._queue.empty(), 'Queue is not clean before next initialization.'
+            self._activated = False
+            self._registry.clear()
+        future = FutureResult()
+        self._registry[identifier] = _MasterRegistry(future)
+        return SlavePipe(identifier, self._queue, future)
+
+    def run_master(self, master_msg):
+        """
+        Main entry for the master device in each forward pass.
+        The messages were first collected from each devices (including the master device), and then
+        an callback will be invoked to compute the message to be sent back to each devices
+        (including the master device).
+
+        Args:
+            master_msg: the message that the master want to send to itself. This will be placed as the first
+            message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example.
+
+        Returns: the message to be sent back to the master device.
+
+        """
+        self._activated = True
+
+        intermediates = [(0, master_msg)]
+        for i in range(self.nr_slaves):
+            intermediates.append(self._queue.get())
+
+        results = self._master_callback(intermediates)
+        assert results[0][0] == 0, 'The first result should belongs to the master.'
+
+        for i, res in results:
+            if i == 0:
+                continue
+            self._registry[i].result.put(res)
+
+        for i in range(self.nr_slaves):
+            assert self._queue.get() is True
+
+        return results[0][1]
+
+    @property
+    def nr_slaves(self):
+        return len(self._registry)

Face_Enhancement/models/networks/Synchronized-BatchNorm-PyTorch/sync_batchnorm/replicate.py

@@ -0,0 +1,94 @@
+# -*- coding: utf-8 -*-
+# File   : replicate.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import functools
+
+from torch.nn.parallel.data_parallel import DataParallel
+
+__all__ = [
+    'CallbackContext',
+    'execute_replication_callbacks',
+    'DataParallelWithCallback',
+    'patch_replication_callback'
+]
+
+
+class CallbackContext(object):
+    pass
+
+
+def execute_replication_callbacks(modules):
+    """
+    Execute an replication callback `__data_parallel_replicate__` on each module created by original replication.
+
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+
+    Note that, as all modules are isomorphism, we assign each sub-module with a context
+    (shared among multiple copies of this module on different devices).
+    Through this context, different copies can share some information.
+
+    We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback
+    of any slave copies.
+    """
+    master_copy = modules[0]
+    nr_modules = len(list(master_copy.modules()))
+    ctxs = [CallbackContext() for _ in range(nr_modules)]
+
+    for i, module in enumerate(modules):
+        for j, m in enumerate(module.modules()):
+            if hasattr(m, '__data_parallel_replicate__'):
+                m.__data_parallel_replicate__(ctxs[j], i)
+
+
+class DataParallelWithCallback(DataParallel):
+    """
+    Data Parallel with a replication callback.
+
+    An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by
+    original `replicate` function.
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+        # sync_bn.__data_parallel_replicate__ will be invoked.
+    """
+
+    def replicate(self, module, device_ids):
+        modules = super(DataParallelWithCallback, self).replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+
+def patch_replication_callback(data_parallel):
+    """
+    Monkey-patch an existing `DataParallel` object. Add the replication callback.
+    Useful when you have customized `DataParallel` implementation.
+
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallel(sync_bn, device_ids=[0, 1])
+        > patch_replication_callback(sync_bn)
+        # this is equivalent to
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+    """
+
+    assert isinstance(data_parallel, DataParallel)
+
+    old_replicate = data_parallel.replicate
+
+    @functools.wraps(old_replicate)
+    def new_replicate(module, device_ids):
+        modules = old_replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+    data_parallel.replicate = new_replicate

Face_Enhancement/models/networks/Synchronized-BatchNorm-PyTorch/sync_batchnorm/unittest.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+# File   : unittest.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import unittest
+import torch
+
+
+class TorchTestCase(unittest.TestCase):
+    def assertTensorClose(self, x, y):
+        adiff = float((x - y).abs().max())
+        if (y == 0).all():
+            rdiff = 'NaN'
+        else:
+            rdiff = float((adiff / y).abs().max())
+
+        message = (
+            'Tensor close check failed\n'
+            'adiff={}\n'
+            'rdiff={}\n'
+        ).format(adiff, rdiff)
+        self.assertTrue(torch.allclose(x, y, atol=1e-5, rtol=1e-3), message)
+

Face_Enhancement/models/networks/Synchronized-BatchNorm-PyTorch/tests/test_numeric_batchnorm.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+# File   : test_numeric_batchnorm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+
+import unittest
+
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+
+from sync_batchnorm.unittest import TorchTestCase
+
+
+def handy_var(a, unbias=True):
+    n = a.size(0)
+    asum = a.sum(dim=0)
+    as_sum = (a ** 2).sum(dim=0)  # a square sum
+    sumvar = as_sum - asum * asum / n
+    if unbias:
+        return sumvar / (n - 1)
+    else:
+        return sumvar / n
+
+
+class NumericTestCase(TorchTestCase):
+    def testNumericBatchNorm(self):
+        a = torch.rand(16, 10)
+        bn = nn.BatchNorm1d(10, momentum=1, eps=1e-5, affine=False)
+        bn.train()
+
+        a_var1 = Variable(a, requires_grad=True)
+        b_var1 = bn(a_var1)
+        loss1 = b_var1.sum()
+        loss1.backward()
+
+        a_var2 = Variable(a, requires_grad=True)
+        a_mean2 = a_var2.mean(dim=0, keepdim=True)
+        a_std2 = torch.sqrt(handy_var(a_var2, unbias=False).clamp(min=1e-5))
+        # a_std2 = torch.sqrt(a_var2.var(dim=0, keepdim=True, unbiased=False) + 1e-5)
+        b_var2 = (a_var2 - a_mean2) / a_std2
+        loss2 = b_var2.sum()
+        loss2.backward()
+
+        self.assertTensorClose(bn.running_mean, a.mean(dim=0))
+        self.assertTensorClose(bn.running_var, handy_var(a))
+        self.assertTensorClose(a_var1.data, a_var2.data)
+        self.assertTensorClose(b_var1.data, b_var2.data)
+        self.assertTensorClose(a_var1.grad, a_var2.grad)
+
+
+if __name__ == '__main__':
+    unittest.main()

Face_Enhancement/models/networks/Synchronized-BatchNorm-PyTorch/tests/test_numeric_batchnorm_v2.py

@@ -0,0 +1,62 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : test_numeric_batchnorm_v2.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 11/01/2018
+#
+# Distributed under terms of the MIT license.
+
+"""
+Test the numerical implementation of batch normalization.
+
+Author: acgtyrant.
+See also: https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/14
+"""
+
+import unittest
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+from sync_batchnorm.unittest import TorchTestCase
+from sync_batchnorm.batchnorm_reimpl import BatchNorm2dReimpl
+
+
+class NumericTestCasev2(TorchTestCase):
+    def testNumericBatchNorm(self):
+        CHANNELS = 16
+        batchnorm1 = nn.BatchNorm2d(CHANNELS, momentum=1)
+        optimizer1 = optim.SGD(batchnorm1.parameters(), lr=0.01)
+
+        batchnorm2 = BatchNorm2dReimpl(CHANNELS, momentum=1)
+        batchnorm2.weight.data.copy_(batchnorm1.weight.data)
+        batchnorm2.bias.data.copy_(batchnorm1.bias.data)
+        optimizer2 = optim.SGD(batchnorm2.parameters(), lr=0.01)
+
+        for _ in range(100):
+            input_ = torch.rand(16, CHANNELS, 16, 16)
+
+            input1 = input_.clone().requires_grad_(True)
+            output1 = batchnorm1(input1)
+            output1.sum().backward()
+            optimizer1.step()
+
+            input2 = input_.clone().requires_grad_(True)
+            output2 = batchnorm2(input2)
+            output2.sum().backward()
+            optimizer2.step()
+
+        self.assertTensorClose(input1, input2)
+        self.assertTensorClose(output1, output2)
+        self.assertTensorClose(input1.grad, input2.grad)
+        self.assertTensorClose(batchnorm1.weight.grad, batchnorm2.weight.grad)
+        self.assertTensorClose(batchnorm1.bias.grad, batchnorm2.bias.grad)
+        self.assertTensorClose(batchnorm1.running_mean, batchnorm2.running_mean)
+        self.assertTensorClose(batchnorm2.running_mean, batchnorm2.running_mean)
+
+
+if __name__ == '__main__':
+    unittest.main()
+

Face_Enhancement/models/networks/Synchronized-BatchNorm-PyTorch/tests/test_sync_batchnorm.py

@@ -0,0 +1,114 @@
+# -*- coding: utf-8 -*-
+# File   : test_sync_batchnorm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+
+import unittest
+
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+
+from sync_batchnorm import set_sbn_eps_mode
+from sync_batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, DataParallelWithCallback
+from sync_batchnorm.unittest import TorchTestCase
+
+set_sbn_eps_mode('plus')
+
+
+def handy_var(a, unbias=True):
+    n = a.size(0)
+    asum = a.sum(dim=0)
+    as_sum = (a ** 2).sum(dim=0)  # a square sum
+    sumvar = as_sum - asum * asum / n
+    if unbias:
+        return sumvar / (n - 1)
+    else:
+        return sumvar / n
+
+
+def _find_bn(module):
+    for m in module.modules():
+        if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, SynchronizedBatchNorm1d, SynchronizedBatchNorm2d)):
+            return m
+
+
+class SyncTestCase(TorchTestCase):
+    def _syncParameters(self, bn1, bn2):
+        bn1.reset_parameters()
+        bn2.reset_parameters()
+        if bn1.affine and bn2.affine:
+            bn2.weight.data.copy_(bn1.weight.data)
+            bn2.bias.data.copy_(bn1.bias.data)
+
+    def _checkBatchNormResult(self, bn1, bn2, input, is_train, cuda=False):
+        """Check the forward and backward for the customized batch normalization."""
+        bn1.train(mode=is_train)
+        bn2.train(mode=is_train)
+
+        if cuda:
+            input = input.cuda()
+
+        self._syncParameters(_find_bn(bn1), _find_bn(bn2))
+
+        input1 = Variable(input, requires_grad=True)
+        output1 = bn1(input1)
+        output1.sum().backward()
+        input2 = Variable(input, requires_grad=True)
+        output2 = bn2(input2)
+        output2.sum().backward()
+
+        self.assertTensorClose(input1.data, input2.data)
+        self.assertTensorClose(output1.data, output2.data)
+        self.assertTensorClose(input1.grad, input2.grad)
+        self.assertTensorClose(_find_bn(bn1).running_mean, _find_bn(bn2).running_mean)
+        self.assertTensorClose(_find_bn(bn1).running_var, _find_bn(bn2).running_var)
+
+    def testSyncBatchNormNormalTrain(self):
+        bn = nn.BatchNorm1d(10)
+        sync_bn = SynchronizedBatchNorm1d(10)
+
+        self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10), True)
+
+    def testSyncBatchNormNormalEval(self):
+        bn = nn.BatchNorm1d(10)
+        sync_bn = SynchronizedBatchNorm1d(10)
+
+        self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10), False)
+
+    def testSyncBatchNormSyncTrain(self):
+        bn = nn.BatchNorm1d(10, eps=1e-5, affine=False)
+        sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+
+        bn.cuda()
+        sync_bn.cuda()
+
+        self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10), True, cuda=True)
+
+    def testSyncBatchNormSyncEval(self):
+        bn = nn.BatchNorm1d(10, eps=1e-5, affine=False)
+        sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+
+        bn.cuda()
+        sync_bn.cuda()
+
+        self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10), False, cuda=True)
+
+    def testSyncBatchNorm2DSyncTrain(self):
+        bn = nn.BatchNorm2d(10)
+        sync_bn = SynchronizedBatchNorm2d(10)
+        sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+
+        bn.cuda()
+        sync_bn.cuda()
+
+        self._checkBatchNormResult(bn, sync_bn, torch.rand(16, 10, 16, 16), True, cuda=True)
+
+
+if __name__ == '__main__':
+    unittest.main()

Face_Enhancement/models/networks/__init__.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+from models.networks.base_network import BaseNetwork
+from models.networks.generator import *
+from models.networks.encoder import *
+import util.util as util
+
+
+def find_network_using_name(target_network_name, filename):
+    target_class_name = target_network_name + filename
+    module_name = "models.networks." + filename
+    network = util.find_class_in_module(target_class_name, module_name)
+
+    assert issubclass(network, BaseNetwork), "Class %s should be a subclass of BaseNetwork" % network
+
+    return network
+
+
+def modify_commandline_options(parser, is_train):
+    opt, _ = parser.parse_known_args()
+
+    netG_cls = find_network_using_name(opt.netG, "generator")
+    parser = netG_cls.modify_commandline_options(parser, is_train)
+    if is_train:
+        netD_cls = find_network_using_name(opt.netD, "discriminator")
+        parser = netD_cls.modify_commandline_options(parser, is_train)
+    netE_cls = find_network_using_name("conv", "encoder")
+    parser = netE_cls.modify_commandline_options(parser, is_train)
+
+    return parser
+
+
+def create_network(cls, opt):
+    net = cls(opt)
+    net.print_network()
+    if len(opt.gpu_ids) > 0:
+        assert torch.cuda.is_available()
+        net.cuda()
+    net.init_weights(opt.init_type, opt.init_variance)
+    return net
+
+
+def define_G(opt):
+    netG_cls = find_network_using_name(opt.netG, "generator")
+    return create_network(netG_cls, opt)
+
+
+def define_D(opt):
+    netD_cls = find_network_using_name(opt.netD, "discriminator")
+    return create_network(netD_cls, opt)
+
+
+def define_E(opt):
+    # there exists only one encoder type
+    netE_cls = find_network_using_name("conv", "encoder")
+    return create_network(netE_cls, opt)

Face_Enhancement/models/networks/architecture.py

@@ -0,0 +1,173 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import torch.nn.utils.spectral_norm as spectral_norm
+from models.networks.normalization import SPADE
+
+
+# ResNet block that uses SPADE.
+# It differs from the ResNet block of pix2pixHD in that
+# it takes in the segmentation map as input, learns the skip connection if necessary,
+# and applies normalization first and then convolution.
+# This architecture seemed like a standard architecture for unconditional or
+# class-conditional GAN architecture using residual block.
+# The code was inspired from https://github.com/LMescheder/GAN_stability.
+class SPADEResnetBlock(nn.Module):
+    def __init__(self, fin, fout, opt):
+        super().__init__()
+        # Attributes
+        self.learned_shortcut = fin != fout
+        fmiddle = min(fin, fout)
+
+        self.opt = opt
+        # create conv layers
+        self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1)
+        self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1)
+        if self.learned_shortcut:
+            self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)
+
+        # apply spectral norm if specified
+        if "spectral" in opt.norm_G:
+            self.conv_0 = spectral_norm(self.conv_0)
+            self.conv_1 = spectral_norm(self.conv_1)
+            if self.learned_shortcut:
+                self.conv_s = spectral_norm(self.conv_s)
+
+        # define normalization layers
+        spade_config_str = opt.norm_G.replace("spectral", "")
+        self.norm_0 = SPADE(spade_config_str, fin, opt.semantic_nc, opt)
+        self.norm_1 = SPADE(spade_config_str, fmiddle, opt.semantic_nc, opt)
+        if self.learned_shortcut:
+            self.norm_s = SPADE(spade_config_str, fin, opt.semantic_nc, opt)
+
+    # note the resnet block with SPADE also takes in |seg|,
+    # the semantic segmentation map as input
+    def forward(self, x, seg, degraded_image):
+        x_s = self.shortcut(x, seg, degraded_image)
+
+        dx = self.conv_0(self.actvn(self.norm_0(x, seg, degraded_image)))
+        dx = self.conv_1(self.actvn(self.norm_1(dx, seg, degraded_image)))
+
+        out = x_s + dx
+
+        return out
+
+    def shortcut(self, x, seg, degraded_image):
+        if self.learned_shortcut:
+            x_s = self.conv_s(self.norm_s(x, seg, degraded_image))
+        else:
+            x_s = x
+        return x_s
+
+    def actvn(self, x):
+        return F.leaky_relu(x, 2e-1)
+
+
+# ResNet block used in pix2pixHD
+# We keep the same architecture as pix2pixHD.
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, norm_layer, activation=nn.ReLU(False), kernel_size=3):
+        super().__init__()
+
+        pw = (kernel_size - 1) // 2
+        self.conv_block = nn.Sequential(
+            nn.ReflectionPad2d(pw),
+            norm_layer(nn.Conv2d(dim, dim, kernel_size=kernel_size)),
+            activation,
+            nn.ReflectionPad2d(pw),
+            norm_layer(nn.Conv2d(dim, dim, kernel_size=kernel_size)),
+        )
+
+    def forward(self, x):
+        y = self.conv_block(x)
+        out = x + y
+        return out
+
+
+# VGG architecter, used for the perceptual loss using a pretrained VGG network
+class VGG19(torch.nn.Module):
+    def __init__(self, requires_grad=False):
+        super().__init__()
+        vgg_pretrained_features = torchvision.models.vgg19(pretrained=True).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        for x in range(2):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(2, 7):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(7, 12):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(12, 21):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(21, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h_relu1 = self.slice1(X)
+        h_relu2 = self.slice2(h_relu1)
+        h_relu3 = self.slice3(h_relu2)
+        h_relu4 = self.slice4(h_relu3)
+        h_relu5 = self.slice5(h_relu4)
+        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
+        return out
+
+
+class SPADEResnetBlock_non_spade(nn.Module):
+    def __init__(self, fin, fout, opt):
+        super().__init__()
+        # Attributes
+        self.learned_shortcut = fin != fout
+        fmiddle = min(fin, fout)
+
+        self.opt = opt
+        # create conv layers
+        self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1)
+        self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1)
+        if self.learned_shortcut:
+            self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)
+
+        # apply spectral norm if specified
+        if "spectral" in opt.norm_G:
+            self.conv_0 = spectral_norm(self.conv_0)
+            self.conv_1 = spectral_norm(self.conv_1)
+            if self.learned_shortcut:
+                self.conv_s = spectral_norm(self.conv_s)
+
+        # define normalization layers
+        spade_config_str = opt.norm_G.replace("spectral", "")
+        self.norm_0 = SPADE(spade_config_str, fin, opt.semantic_nc, opt)
+        self.norm_1 = SPADE(spade_config_str, fmiddle, opt.semantic_nc, opt)
+        if self.learned_shortcut:
+            self.norm_s = SPADE(spade_config_str, fin, opt.semantic_nc, opt)
+
+    # note the resnet block with SPADE also takes in |seg|,
+    # the semantic segmentation map as input
+    def forward(self, x, seg, degraded_image):
+        x_s = self.shortcut(x, seg, degraded_image)
+
+        dx = self.conv_0(self.actvn(x))
+        dx = self.conv_1(self.actvn(dx))
+
+        out = x_s + dx
+
+        return out
+
+    def shortcut(self, x, seg, degraded_image):
+        if self.learned_shortcut:
+            x_s = self.conv_s(x)
+        else:
+            x_s = x
+        return x_s
+
+    def actvn(self, x):
+        return F.leaky_relu(x, 2e-1)

Face_Enhancement/models/networks/base_network.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch.nn as nn
+from torch.nn import init
+
+
+class BaseNetwork(nn.Module):
+    def __init__(self):
+        super(BaseNetwork, self).__init__()
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        return parser
+
+    def print_network(self):
+        if isinstance(self, list):
+            self = self[0]
+        num_params = 0
+        for param in self.parameters():
+            num_params += param.numel()
+        print(
+            "Network [%s] was created. Total number of parameters: %.1f million. "
+            "To see the architecture, do print(network)." % (type(self).__name__, num_params / 1000000)
+        )
+
+    def init_weights(self, init_type="normal", gain=0.02):
+        def init_func(m):
+            classname = m.__class__.__name__
+            if classname.find("BatchNorm2d") != -1:
+                if hasattr(m, "weight") and m.weight is not None:
+                    init.normal_(m.weight.data, 1.0, gain)
+                if hasattr(m, "bias") and m.bias is not None:
+                    init.constant_(m.bias.data, 0.0)
+            elif hasattr(m, "weight") and (classname.find("Conv") != -1 or classname.find("Linear") != -1):
+                if init_type == "normal":
+                    init.normal_(m.weight.data, 0.0, gain)
+                elif init_type == "xavier":
+                    init.xavier_normal_(m.weight.data, gain=gain)
+                elif init_type == "xavier_uniform":
+                    init.xavier_uniform_(m.weight.data, gain=1.0)
+                elif init_type == "kaiming":
+                    init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
+                elif init_type == "orthogonal":
+                    init.orthogonal_(m.weight.data, 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:
+                    init.constant_(m.bias.data, 0.0)
+
+        self.apply(init_func)
+
+        # propagate to children
+        for m in self.children():
+            if hasattr(m, "init_weights"):
+                m.init_weights(init_type, gain)

Face_Enhancement/models/networks/encoder.py

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+from models.networks.base_network import BaseNetwork
+from models.networks.normalization import get_nonspade_norm_layer
+
+
+class ConvEncoder(BaseNetwork):
+    """ Same architecture as the image discriminator """
+
+    def __init__(self, opt):
+        super().__init__()
+
+        kw = 3
+        pw = int(np.ceil((kw - 1.0) / 2))
+        ndf = opt.ngf
+        norm_layer = get_nonspade_norm_layer(opt, opt.norm_E)
+        self.layer1 = norm_layer(nn.Conv2d(3, ndf, kw, stride=2, padding=pw))
+        self.layer2 = norm_layer(nn.Conv2d(ndf * 1, ndf * 2, kw, stride=2, padding=pw))
+        self.layer3 = norm_layer(nn.Conv2d(ndf * 2, ndf * 4, kw, stride=2, padding=pw))
+        self.layer4 = norm_layer(nn.Conv2d(ndf * 4, ndf * 8, kw, stride=2, padding=pw))
+        self.layer5 = norm_layer(nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw))
+        if opt.crop_size >= 256:
+            self.layer6 = norm_layer(nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw))
+
+        self.so = s0 = 4
+        self.fc_mu = nn.Linear(ndf * 8 * s0 * s0, 256)
+        self.fc_var = nn.Linear(ndf * 8 * s0 * s0, 256)
+
+        self.actvn = nn.LeakyReLU(0.2, False)
+        self.opt = opt
+
+    def forward(self, x):
+        if x.size(2) != 256 or x.size(3) != 256:
+            x = F.interpolate(x, size=(256, 256), mode="bilinear")
+
+        x = self.layer1(x)
+        x = self.layer2(self.actvn(x))
+        x = self.layer3(self.actvn(x))
+        x = self.layer4(self.actvn(x))
+        x = self.layer5(self.actvn(x))
+        if self.opt.crop_size >= 256:
+            x = self.layer6(self.actvn(x))
+        x = self.actvn(x)
+
+        x = x.view(x.size(0), -1)
+        mu = self.fc_mu(x)
+        logvar = self.fc_var(x)
+
+        return mu, logvar

Face_Enhancement/models/networks/generator.py

@@ -0,0 +1,233 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models.networks.base_network import BaseNetwork
+from models.networks.normalization import get_nonspade_norm_layer
+from models.networks.architecture import ResnetBlock as ResnetBlock
+from models.networks.architecture import SPADEResnetBlock as SPADEResnetBlock
+from models.networks.architecture import SPADEResnetBlock_non_spade as SPADEResnetBlock_non_spade
+
+
+class SPADEGenerator(BaseNetwork):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        parser.set_defaults(norm_G="spectralspadesyncbatch3x3")
+        parser.add_argument(
+            "--num_upsampling_layers",
+            choices=("normal", "more", "most"),
+            default="normal",
+            help="If 'more', adds upsampling layer between the two middle resnet blocks. If 'most', also add one more upsampling + resnet layer at the end of the generator",
+        )
+
+        return parser
+
+    def __init__(self, opt):
+        super().__init__()
+        self.opt = opt
+        nf = opt.ngf
+
+        self.sw, self.sh = self.compute_latent_vector_size(opt)
+
+        print("The size of the latent vector size is [%d,%d]" % (self.sw, self.sh))
+
+        if opt.use_vae:
+            # In case of VAE, we will sample from random z vector
+            self.fc = nn.Linear(opt.z_dim, 16 * nf * self.sw * self.sh)
+        else:
+            # Otherwise, we make the network deterministic by starting with
+            # downsampled segmentation map instead of random z
+            if self.opt.no_parsing_map:
+                self.fc = nn.Conv2d(3, 16 * nf, 3, padding=1)
+            else:
+                self.fc = nn.Conv2d(self.opt.semantic_nc, 16 * nf, 3, padding=1)
+
+        if self.opt.injection_layer == "all" or self.opt.injection_layer == "1":
+            self.head_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt)
+        else:
+            self.head_0 = SPADEResnetBlock_non_spade(16 * nf, 16 * nf, opt)
+
+        if self.opt.injection_layer == "all" or self.opt.injection_layer == "2":
+            self.G_middle_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt)
+            self.G_middle_1 = SPADEResnetBlock(16 * nf, 16 * nf, opt)
+
+        else:
+            self.G_middle_0 = SPADEResnetBlock_non_spade(16 * nf, 16 * nf, opt)
+            self.G_middle_1 = SPADEResnetBlock_non_spade(16 * nf, 16 * nf, opt)
+
+        if self.opt.injection_layer == "all" or self.opt.injection_layer == "3":
+            self.up_0 = SPADEResnetBlock(16 * nf, 8 * nf, opt)
+        else:
+            self.up_0 = SPADEResnetBlock_non_spade(16 * nf, 8 * nf, opt)
+
+        if self.opt.injection_layer == "all" or self.opt.injection_layer == "4":
+            self.up_1 = SPADEResnetBlock(8 * nf, 4 * nf, opt)
+        else:
+            self.up_1 = SPADEResnetBlock_non_spade(8 * nf, 4 * nf, opt)
+
+        if self.opt.injection_layer == "all" or self.opt.injection_layer == "5":
+            self.up_2 = SPADEResnetBlock(4 * nf, 2 * nf, opt)
+        else:
+            self.up_2 = SPADEResnetBlock_non_spade(4 * nf, 2 * nf, opt)
+
+        if self.opt.injection_layer == "all" or self.opt.injection_layer == "6":
+            self.up_3 = SPADEResnetBlock(2 * nf, 1 * nf, opt)
+        else:
+            self.up_3 = SPADEResnetBlock_non_spade(2 * nf, 1 * nf, opt)
+
+        final_nc = nf
+
+        if opt.num_upsampling_layers == "most":
+            self.up_4 = SPADEResnetBlock(1 * nf, nf // 2, opt)
+            final_nc = nf // 2
+
+        self.conv_img = nn.Conv2d(final_nc, 3, 3, padding=1)
+
+        self.up = nn.Upsample(scale_factor=2)
+
+    def compute_latent_vector_size(self, opt):
+        if opt.num_upsampling_layers == "normal":
+            num_up_layers = 5
+        elif opt.num_upsampling_layers == "more":
+            num_up_layers = 6
+        elif opt.num_upsampling_layers == "most":
+            num_up_layers = 7
+        else:
+            raise ValueError("opt.num_upsampling_layers [%s] not recognized" % opt.num_upsampling_layers)
+
+        sw = opt.load_size // (2 ** num_up_layers)
+        sh = round(sw / opt.aspect_ratio)
+
+        return sw, sh
+
+    def forward(self, input, degraded_image, z=None):
+        seg = input
+
+        if self.opt.use_vae:
+            # we sample z from unit normal and reshape the tensor
+            if z is None:
+                z = torch.randn(input.size(0), self.opt.z_dim, dtype=torch.float32, device=input.get_device())
+            x = self.fc(z)
+            x = x.view(-1, 16 * self.opt.ngf, self.sh, self.sw)
+        else:
+            # we downsample segmap and run convolution
+            if self.opt.no_parsing_map:
+                x = F.interpolate(degraded_image, size=(self.sh, self.sw), mode="bilinear")
+            else:
+                x = F.interpolate(seg, size=(self.sh, self.sw), mode="nearest")
+            x = self.fc(x)
+
+        x = self.head_0(x, seg, degraded_image)
+
+        x = self.up(x)
+        x = self.G_middle_0(x, seg, degraded_image)
+
+        if self.opt.num_upsampling_layers == "more" or self.opt.num_upsampling_layers == "most":
+            x = self.up(x)
+
+        x = self.G_middle_1(x, seg, degraded_image)
+
+        x = self.up(x)
+        x = self.up_0(x, seg, degraded_image)
+        x = self.up(x)
+        x = self.up_1(x, seg, degraded_image)
+        x = self.up(x)
+        x = self.up_2(x, seg, degraded_image)
+        x = self.up(x)
+        x = self.up_3(x, seg, degraded_image)
+
+        if self.opt.num_upsampling_layers == "most":
+            x = self.up(x)
+            x = self.up_4(x, seg, degraded_image)
+
+        x = self.conv_img(F.leaky_relu(x, 2e-1))
+        x = F.tanh(x)
+
+        return x
+
+
+class Pix2PixHDGenerator(BaseNetwork):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        parser.add_argument(
+            "--resnet_n_downsample", type=int, default=4, help="number of downsampling layers in netG"
+        )
+        parser.add_argument(
+            "--resnet_n_blocks",
+            type=int,
+            default=9,
+            help="number of residual blocks in the global generator network",
+        )
+        parser.add_argument(
+            "--resnet_kernel_size", type=int, default=3, help="kernel size of the resnet block"
+        )
+        parser.add_argument(
+            "--resnet_initial_kernel_size", type=int, default=7, help="kernel size of the first convolution"
+        )
+        # parser.set_defaults(norm_G='instance')
+        return parser
+
+    def __init__(self, opt):
+        super().__init__()
+        input_nc = 3
+
+        # print("xxxxx")
+        # print(opt.norm_G)
+        norm_layer = get_nonspade_norm_layer(opt, opt.norm_G)
+        activation = nn.ReLU(False)
+
+        model = []
+
+        # initial conv
+        model += [
+            nn.ReflectionPad2d(opt.resnet_initial_kernel_size // 2),
+            norm_layer(nn.Conv2d(input_nc, opt.ngf, kernel_size=opt.resnet_initial_kernel_size, padding=0)),
+            activation,
+        ]
+
+        # downsample
+        mult = 1
+        for i in range(opt.resnet_n_downsample):
+            model += [
+                norm_layer(nn.Conv2d(opt.ngf * mult, opt.ngf * mult * 2, kernel_size=3, stride=2, padding=1)),
+                activation,
+            ]
+            mult *= 2
+
+        # resnet blocks
+        for i in range(opt.resnet_n_blocks):
+            model += [
+                ResnetBlock(
+                    opt.ngf * mult,
+                    norm_layer=norm_layer,
+                    activation=activation,
+                    kernel_size=opt.resnet_kernel_size,
+                )
+            ]
+
+        # upsample
+        for i in range(opt.resnet_n_downsample):
+            nc_in = int(opt.ngf * mult)
+            nc_out = int((opt.ngf * mult) / 2)
+            model += [
+                norm_layer(
+                    nn.ConvTranspose2d(nc_in, nc_out, kernel_size=3, stride=2, padding=1, output_padding=1)
+                ),
+                activation,
+            ]
+            mult = mult // 2
+
+        # final output conv
+        model += [
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(nc_out, opt.output_nc, kernel_size=7, padding=0),
+            nn.Tanh(),
+        ]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input, degraded_image, z=None):
+        return self.model(degraded_image)
+

Face_Enhancement/models/networks/normalization.py

@@ -0,0 +1,100 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import re
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models.networks.sync_batchnorm import SynchronizedBatchNorm2d
+import torch.nn.utils.spectral_norm as spectral_norm
+
+
+def get_nonspade_norm_layer(opt, norm_type="instance"):
+    # helper function to get # output channels of the previous layer
+    def get_out_channel(layer):
+        if hasattr(layer, "out_channels"):
+            return getattr(layer, "out_channels")
+        return layer.weight.size(0)
+
+    # this function will be returned
+    def add_norm_layer(layer):
+        nonlocal norm_type
+        if norm_type.startswith("spectral"):
+            layer = spectral_norm(layer)
+            subnorm_type = norm_type[len("spectral") :]
+
+        if subnorm_type == "none" or len(subnorm_type) == 0:
+            return layer
+
+        # remove bias in the previous layer, which is meaningless
+        # since it has no effect after normalization
+        if getattr(layer, "bias", None) is not None:
+            delattr(layer, "bias")
+            layer.register_parameter("bias", None)
+
+        if subnorm_type == "batch":
+            norm_layer = nn.BatchNorm2d(get_out_channel(layer), affine=True)
+        elif subnorm_type == "sync_batch":
+            norm_layer = SynchronizedBatchNorm2d(get_out_channel(layer), affine=True)
+        elif subnorm_type == "instance":
+            norm_layer = nn.InstanceNorm2d(get_out_channel(layer), affine=False)
+        else:
+            raise ValueError("normalization layer %s is not recognized" % subnorm_type)
+
+        return nn.Sequential(layer, norm_layer)
+
+    return add_norm_layer
+
+
+class SPADE(nn.Module):
+    def __init__(self, config_text, norm_nc, label_nc, opt):
+        super().__init__()
+
+        assert config_text.startswith("spade")
+        parsed = re.search("spade(\D+)(\d)x\d", config_text)
+        param_free_norm_type = str(parsed.group(1))
+        ks = int(parsed.group(2))
+        self.opt = opt
+        if param_free_norm_type == "instance":
+            self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
+        elif param_free_norm_type == "syncbatch":
+            self.param_free_norm = SynchronizedBatchNorm2d(norm_nc, affine=False)
+        elif param_free_norm_type == "batch":
+            self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False)
+        else:
+            raise ValueError("%s is not a recognized param-free norm type in SPADE" % param_free_norm_type)
+
+        # The dimension of the intermediate embedding space. Yes, hardcoded.
+        nhidden = 128
+
+        pw = ks // 2
+
+        if self.opt.no_parsing_map:
+            self.mlp_shared = nn.Sequential(nn.Conv2d(3, nhidden, kernel_size=ks, padding=pw), nn.ReLU())
+        else:
+            self.mlp_shared = nn.Sequential(
+                nn.Conv2d(label_nc + 3, nhidden, kernel_size=ks, padding=pw), nn.ReLU()
+            )
+        self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)
+        self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)
+
+    def forward(self, x, segmap, degraded_image):
+
+        # Part 1. generate parameter-free normalized activations
+        normalized = self.param_free_norm(x)
+
+        # Part 2. produce scaling and bias conditioned on semantic map
+        segmap = F.interpolate(segmap, size=x.size()[2:], mode="nearest")
+        degraded_face = F.interpolate(degraded_image, size=x.size()[2:], mode="bilinear")
+
+        if self.opt.no_parsing_map:
+            actv = self.mlp_shared(degraded_face)
+        else:
+            actv = self.mlp_shared(torch.cat((segmap, degraded_face), dim=1))
+        gamma = self.mlp_gamma(actv)
+        beta = self.mlp_beta(actv)
+
+        # apply scale and bias
+        out = normalized * (1 + gamma) + beta
+
+        return out

Face_Enhancement/models/networks/sync_batchnorm/__init__.py

@@ -0,0 +1,14 @@
+# -*- coding: utf-8 -*-
+# File   : __init__.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+from .batchnorm import set_sbn_eps_mode
+from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d
+from .batchnorm import patch_sync_batchnorm, convert_model
+from .replicate import DataParallelWithCallback, patch_replication_callback

Face_Enhancement/models/networks/sync_batchnorm/batchnorm.py

@@ -0,0 +1,412 @@
+# -*- coding: utf-8 -*-
+# File   : batchnorm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import collections
+import contextlib
+
+import torch
+import torch.nn.functional as F
+
+from torch.nn.modules.batchnorm import _BatchNorm
+
+try:
+    from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast
+except ImportError:
+    ReduceAddCoalesced = Broadcast = None
+
+try:
+    from jactorch.parallel.comm import SyncMaster
+    from jactorch.parallel.data_parallel import JacDataParallel as DataParallelWithCallback
+except ImportError:
+    from .comm import SyncMaster
+    from .replicate import DataParallelWithCallback
+
+__all__ = [
+    'set_sbn_eps_mode',
+    'SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d', 'SynchronizedBatchNorm3d',
+    'patch_sync_batchnorm', 'convert_model'
+]
+
+
+SBN_EPS_MODE = 'clamp'
+
+
+def set_sbn_eps_mode(mode):
+    global SBN_EPS_MODE
+    assert mode in ('clamp', 'plus')
+    SBN_EPS_MODE = mode
+
+
+def _sum_ft(tensor):
+    """sum over the first and last dimention"""
+    return tensor.sum(dim=0).sum(dim=-1)
+
+
+def _unsqueeze_ft(tensor):
+    """add new dimensions at the front and the tail"""
+    return tensor.unsqueeze(0).unsqueeze(-1)
+
+
+_ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size'])
+_MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std'])
+
+
+class _SynchronizedBatchNorm(_BatchNorm):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True):
+        assert ReduceAddCoalesced is not None, 'Can not use Synchronized Batch Normalization without CUDA support.'
+
+        super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine,
+                                                     track_running_stats=track_running_stats)
+
+        if not self.track_running_stats:
+            import warnings
+            warnings.warn('track_running_stats=False is not supported by the SynchronizedBatchNorm.')
+
+        self._sync_master = SyncMaster(self._data_parallel_master)
+
+        self._is_parallel = False
+        self._parallel_id = None
+        self._slave_pipe = None
+
+    def forward(self, input):
+        # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation.
+        if not (self._is_parallel and self.training):
+            return F.batch_norm(
+                input, self.running_mean, self.running_var, self.weight, self.bias,
+                self.training, self.momentum, self.eps)
+
+        # Resize the input to (B, C, -1).
+        input_shape = input.size()
+        assert input.size(1) == self.num_features, 'Channel size mismatch: got {}, expect {}.'.format(input.size(1), self.num_features)
+        input = input.view(input.size(0), self.num_features, -1)
+
+        # Compute the sum and square-sum.
+        sum_size = input.size(0) * input.size(2)
+        input_sum = _sum_ft(input)
+        input_ssum = _sum_ft(input ** 2)
+
+        # Reduce-and-broadcast the statistics.
+        if self._parallel_id == 0:
+            mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size))
+        else:
+            mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size))
+
+        # Compute the output.
+        if self.affine:
+            # MJY:: Fuse the multiplication for speed.
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias)
+        else:
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std)
+
+        # Reshape it.
+        return output.view(input_shape)
+
+    def __data_parallel_replicate__(self, ctx, copy_id):
+        self._is_parallel = True
+        self._parallel_id = copy_id
+
+        # parallel_id == 0 means master device.
+        if self._parallel_id == 0:
+            ctx.sync_master = self._sync_master
+        else:
+            self._slave_pipe = ctx.sync_master.register_slave(copy_id)
+
+    def _data_parallel_master(self, intermediates):
+        """Reduce the sum and square-sum, compute the statistics, and broadcast it."""
+
+        # Always using same "device order" makes the ReduceAdd operation faster.
+        # Thanks to:: Tete Xiao (http://tetexiao.com/)
+        intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device())
+
+        to_reduce = [i[1][:2] for i in intermediates]
+        to_reduce = [j for i in to_reduce for j in i]  # flatten
+        target_gpus = [i[1].sum.get_device() for i in intermediates]
+
+        sum_size = sum([i[1].sum_size for i in intermediates])
+        sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce)
+        mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size)
+
+        broadcasted = Broadcast.apply(target_gpus, mean, inv_std)
+
+        outputs = []
+        for i, rec in enumerate(intermediates):
+            outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2])))
+
+        return outputs
+
+    def _compute_mean_std(self, sum_, ssum, size):
+        """Compute the mean and standard-deviation with sum and square-sum. This method
+        also maintains the moving average on the master device."""
+        assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.'
+        mean = sum_ / size
+        sumvar = ssum - sum_ * mean
+        unbias_var = sumvar / (size - 1)
+        bias_var = sumvar / size
+
+        if hasattr(torch, 'no_grad'):
+            with torch.no_grad():
+                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data
+                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data
+        else:
+            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data
+            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data
+
+        if SBN_EPS_MODE == 'clamp':
+            return mean, bias_var.clamp(self.eps) ** -0.5
+        elif SBN_EPS_MODE == 'plus':
+            return mean, (bias_var + self.eps) ** -0.5
+        else:
+            raise ValueError('Unknown EPS mode: {}.'.format(SBN_EPS_MODE))
+
+
+class SynchronizedBatchNorm1d(_SynchronizedBatchNorm):
+    r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a
+    mini-batch.
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm1d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of size
+            `batch_size x num_features [x width]`
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape::
+        - Input: :math:`(N, C)` or :math:`(N, C, L)`
+        - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 2 and input.dim() != 3:
+            raise ValueError('expected 2D or 3D input (got {}D input)'
+                             .format(input.dim()))
+
+
+class SynchronizedBatchNorm2d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch
+    of 3d inputs
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm2d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape::
+        - Input: :math:`(N, C, H, W)`
+        - Output: :math:`(N, C, H, W)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 4:
+            raise ValueError('expected 4D input (got {}D input)'
+                             .format(input.dim()))
+
+
+class SynchronizedBatchNorm3d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch
+    of 4d inputs
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm3d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm
+    or Spatio-temporal BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x depth x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape::
+        - Input: :math:`(N, C, D, H, W)`
+        - Output: :math:`(N, C, D, H, W)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 5:
+            raise ValueError('expected 5D input (got {}D input)'
+                             .format(input.dim()))
+
+
+@contextlib.contextmanager
+def patch_sync_batchnorm():
+    import torch.nn as nn
+
+    backup = nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d
+
+    nn.BatchNorm1d = SynchronizedBatchNorm1d
+    nn.BatchNorm2d = SynchronizedBatchNorm2d
+    nn.BatchNorm3d = SynchronizedBatchNorm3d
+
+    yield
+
+    nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d = backup
+
+
+def convert_model(module):
+    """Traverse the input module and its child recursively
+       and replace all instance of torch.nn.modules.batchnorm.BatchNorm*N*d
+       to SynchronizedBatchNorm*N*d
+
+    Args:
+        module: the input module needs to be convert to SyncBN model
+
+    Examples:
+        >>> import torch.nn as nn
+        >>> import torchvision
+        >>> # m is a standard pytorch model
+        >>> m = torchvision.models.resnet18(True)
+        >>> m = nn.DataParallel(m)
+        >>> # after convert, m is using SyncBN
+        >>> m = convert_model(m)
+    """
+    if isinstance(module, torch.nn.DataParallel):
+        mod = module.module
+        mod = convert_model(mod)
+        mod = DataParallelWithCallback(mod, device_ids=module.device_ids)
+        return mod
+
+    mod = module
+    for pth_module, sync_module in zip([torch.nn.modules.batchnorm.BatchNorm1d,
+                                        torch.nn.modules.batchnorm.BatchNorm2d,
+                                        torch.nn.modules.batchnorm.BatchNorm3d],
+                                       [SynchronizedBatchNorm1d,
+                                        SynchronizedBatchNorm2d,
+                                        SynchronizedBatchNorm3d]):
+        if isinstance(module, pth_module):
+            mod = sync_module(module.num_features, module.eps, module.momentum, module.affine)
+            mod.running_mean = module.running_mean
+            mod.running_var = module.running_var
+            if module.affine:
+                mod.weight.data = module.weight.data.clone().detach()
+                mod.bias.data = module.bias.data.clone().detach()
+
+    for name, child in module.named_children():
+        mod.add_module(name, convert_model(child))
+
+    return mod

Face_Enhancement/models/networks/sync_batchnorm/batchnorm_reimpl.py

@@ -0,0 +1,74 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : batchnorm_reimpl.py
+# Author : acgtyrant
+# Date   : 11/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+
+__all__ = ['BatchNorm2dReimpl']
+
+
+class BatchNorm2dReimpl(nn.Module):
+    """
+    A re-implementation of batch normalization, used for testing the numerical
+    stability.
+
+    Author: acgtyrant
+    See also:
+    https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/14
+    """
+    def __init__(self, num_features, eps=1e-5, momentum=0.1):
+        super().__init__()
+
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.weight = nn.Parameter(torch.empty(num_features))
+        self.bias = nn.Parameter(torch.empty(num_features))
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+        self.reset_parameters()
+
+    def reset_running_stats(self):
+        self.running_mean.zero_()
+        self.running_var.fill_(1)
+
+    def reset_parameters(self):
+        self.reset_running_stats()
+        init.uniform_(self.weight)
+        init.zeros_(self.bias)
+
+    def forward(self, input_):
+        batchsize, channels, height, width = input_.size()
+        numel = batchsize * height * width
+        input_ = input_.permute(1, 0, 2, 3).contiguous().view(channels, numel)
+        sum_ = input_.sum(1)
+        sum_of_square = input_.pow(2).sum(1)
+        mean = sum_ / numel
+        sumvar = sum_of_square - sum_ * mean
+
+        self.running_mean = (
+                (1 - self.momentum) * self.running_mean
+                + self.momentum * mean.detach()
+        )
+        unbias_var = sumvar / (numel - 1)
+        self.running_var = (
+                (1 - self.momentum) * self.running_var
+                + self.momentum * unbias_var.detach()
+        )
+
+        bias_var = sumvar / numel
+        inv_std = 1 / (bias_var + self.eps).pow(0.5)
+        output = (
+                (input_ - mean.unsqueeze(1)) * inv_std.unsqueeze(1) *
+                self.weight.unsqueeze(1) + self.bias.unsqueeze(1))
+
+        return output.view(channels, batchsize, height, width).permute(1, 0, 2, 3).contiguous()
+

Face_Enhancement/models/networks/sync_batchnorm/comm.py

@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+# File   : comm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import queue
+import collections
+import threading
+
+__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster']
+
+
+class FutureResult(object):
+    """A thread-safe future implementation. Used only as one-to-one pipe."""
+
+    def __init__(self):
+        self._result = None
+        self._lock = threading.Lock()
+        self._cond = threading.Condition(self._lock)
+
+    def put(self, result):
+        with self._lock:
+            assert self._result is None, 'Previous result has\'t been fetched.'
+            self._result = result
+            self._cond.notify()
+
+    def get(self):
+        with self._lock:
+            if self._result is None:
+                self._cond.wait()
+
+            res = self._result
+            self._result = None
+            return res
+
+
+_MasterRegistry = collections.namedtuple('MasterRegistry', ['result'])
+_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result'])
+
+
+class SlavePipe(_SlavePipeBase):
+    """Pipe for master-slave communication."""
+
+    def run_slave(self, msg):
+        self.queue.put((self.identifier, msg))
+        ret = self.result.get()
+        self.queue.put(True)
+        return ret
+
+
+class SyncMaster(object):
+    """An abstract `SyncMaster` object.
+
+    - During the replication, as the data parallel will trigger an callback of each module, all slave devices should
+    call `register(id)` and obtain an `SlavePipe` to communicate with the master.
+    - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected,
+    and passed to a registered callback.
+    - After receiving the messages, the master device should gather the information and determine to message passed
+    back to each slave devices.
+    """
+
+    def __init__(self, master_callback):
+        """
+
+        Args:
+            master_callback: a callback to be invoked after having collected messages from slave devices.
+        """
+        self._master_callback = master_callback
+        self._queue = queue.Queue()
+        self._registry = collections.OrderedDict()
+        self._activated = False
+
+    def __getstate__(self):
+        return {'master_callback': self._master_callback}
+
+    def __setstate__(self, state):
+        self.__init__(state['master_callback'])
+
+    def register_slave(self, identifier):
+        """
+        Register an slave device.
+
+        Args:
+            identifier: an identifier, usually is the device id.
+
+        Returns: a `SlavePipe` object which can be used to communicate with the master device.
+
+        """
+        if self._activated:
+            assert self._queue.empty(), 'Queue is not clean before next initialization.'
+            self._activated = False
+            self._registry.clear()
+        future = FutureResult()
+        self._registry[identifier] = _MasterRegistry(future)
+        return SlavePipe(identifier, self._queue, future)
+
+    def run_master(self, master_msg):
+        """
+        Main entry for the master device in each forward pass.
+        The messages were first collected from each devices (including the master device), and then
+        an callback will be invoked to compute the message to be sent back to each devices
+        (including the master device).
+
+        Args:
+            master_msg: the message that the master want to send to itself. This will be placed as the first
+            message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example.
+
+        Returns: the message to be sent back to the master device.
+
+        """
+        self._activated = True
+
+        intermediates = [(0, master_msg)]
+        for i in range(self.nr_slaves):
+            intermediates.append(self._queue.get())
+
+        results = self._master_callback(intermediates)
+        assert results[0][0] == 0, 'The first result should belongs to the master.'
+
+        for i, res in results:
+            if i == 0:
+                continue
+            self._registry[i].result.put(res)
+
+        for i in range(self.nr_slaves):
+            assert self._queue.get() is True
+
+        return results[0][1]
+
+    @property
+    def nr_slaves(self):
+        return len(self._registry)

Face_Enhancement/models/networks/sync_batchnorm/replicate.py

@@ -0,0 +1,94 @@
+# -*- coding: utf-8 -*-
+# File   : replicate.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import functools
+
+from torch.nn.parallel.data_parallel import DataParallel
+
+__all__ = [
+    'CallbackContext',
+    'execute_replication_callbacks',
+    'DataParallelWithCallback',
+    'patch_replication_callback'
+]
+
+
+class CallbackContext(object):
+    pass
+
+
+def execute_replication_callbacks(modules):
+    """
+    Execute an replication callback `__data_parallel_replicate__` on each module created by original replication.
+
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+
+    Note that, as all modules are isomorphism, we assign each sub-module with a context
+    (shared among multiple copies of this module on different devices).
+    Through this context, different copies can share some information.
+
+    We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback
+    of any slave copies.
+    """
+    master_copy = modules[0]
+    nr_modules = len(list(master_copy.modules()))
+    ctxs = [CallbackContext() for _ in range(nr_modules)]
+
+    for i, module in enumerate(modules):
+        for j, m in enumerate(module.modules()):
+            if hasattr(m, '__data_parallel_replicate__'):
+                m.__data_parallel_replicate__(ctxs[j], i)
+
+
+class DataParallelWithCallback(DataParallel):
+    """
+    Data Parallel with a replication callback.
+
+    An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by
+    original `replicate` function.
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+        # sync_bn.__data_parallel_replicate__ will be invoked.
+    """
+
+    def replicate(self, module, device_ids):
+        modules = super(DataParallelWithCallback, self).replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+
+def patch_replication_callback(data_parallel):
+    """
+    Monkey-patch an existing `DataParallel` object. Add the replication callback.
+    Useful when you have customized `DataParallel` implementation.
+
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallel(sync_bn, device_ids=[0, 1])
+        > patch_replication_callback(sync_bn)
+        # this is equivalent to
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+    """
+
+    assert isinstance(data_parallel, DataParallel)
+
+    old_replicate = data_parallel.replicate
+
+    @functools.wraps(old_replicate)
+    def new_replicate(module, device_ids):
+        modules = old_replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+    data_parallel.replicate = new_replicate

Face_Enhancement/models/networks/sync_batchnorm/unittest.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+# File   : unittest.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import unittest
+import torch
+
+
+class TorchTestCase(unittest.TestCase):
+    def assertTensorClose(self, x, y):
+        adiff = float((x - y).abs().max())
+        if (y == 0).all():
+            rdiff = 'NaN'
+        else:
+            rdiff = float((adiff / y).abs().max())
+
+        message = (
+            'Tensor close check failed\n'
+            'adiff={}\n'
+            'rdiff={}\n'
+        ).format(adiff, rdiff)
+        self.assertTrue(torch.allclose(x, y, atol=1e-5, rtol=1e-3), message)
+

Face_Enhancement/models/pix2pix_model.py

@@ -0,0 +1,246 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import torch
+import models.networks as networks
+import util.util as util
+
+
+class Pix2PixModel(torch.nn.Module):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        networks.modify_commandline_options(parser, is_train)
+        return parser
+
+    def __init__(self, opt):
+        super().__init__()
+        self.opt = opt
+        self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() else torch.FloatTensor
+        self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() else torch.ByteTensor
+
+        self.netG, self.netD, self.netE = self.initialize_networks(opt)
+
+        # set loss functions
+        if opt.isTrain:
+            self.criterionGAN = networks.GANLoss(opt.gan_mode, tensor=self.FloatTensor, opt=self.opt)
+            self.criterionFeat = torch.nn.L1Loss()
+            if not opt.no_vgg_loss:
+                self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)
+            if opt.use_vae:
+                self.KLDLoss = networks.KLDLoss()
+
+    # Entry point for all calls involving forward pass
+    # of deep networks. We used this approach since DataParallel module
+    # can't parallelize custom functions, we branch to different
+    # routines based on |mode|.
+    def forward(self, data, mode):
+        input_semantics, real_image, degraded_image = self.preprocess_input(data)
+
+        if mode == "generator":
+            g_loss, generated = self.compute_generator_loss(input_semantics, degraded_image, real_image)
+            return g_loss, generated
+        elif mode == "discriminator":
+            d_loss = self.compute_discriminator_loss(input_semantics, degraded_image, real_image)
+            return d_loss
+        elif mode == "encode_only":
+            z, mu, logvar = self.encode_z(real_image)
+            return mu, logvar
+        elif mode == "inference":
+            with torch.no_grad():
+                fake_image, _ = self.generate_fake(input_semantics, degraded_image, real_image)
+            return fake_image
+        else:
+            raise ValueError("|mode| is invalid")
+
+    def create_optimizers(self, opt):
+        G_params = list(self.netG.parameters())
+        if opt.use_vae:
+            G_params += list(self.netE.parameters())
+        if opt.isTrain:
+            D_params = list(self.netD.parameters())
+
+        beta1, beta2 = opt.beta1, opt.beta2
+        if opt.no_TTUR:
+            G_lr, D_lr = opt.lr, opt.lr
+        else:
+            G_lr, D_lr = opt.lr / 2, opt.lr * 2
+
+        optimizer_G = torch.optim.Adam(G_params, lr=G_lr, betas=(beta1, beta2))
+        optimizer_D = torch.optim.Adam(D_params, lr=D_lr, betas=(beta1, beta2))
+
+        return optimizer_G, optimizer_D
+
+    def save(self, epoch):
+        util.save_network(self.netG, "G", epoch, self.opt)
+        util.save_network(self.netD, "D", epoch, self.opt)
+        if self.opt.use_vae:
+            util.save_network(self.netE, "E", epoch, self.opt)
+
+    ############################################################################
+    # Private helper methods
+    ############################################################################
+
+    def initialize_networks(self, opt):
+        netG = networks.define_G(opt)
+        netD = networks.define_D(opt) if opt.isTrain else None
+        netE = networks.define_E(opt) if opt.use_vae else None
+
+        if not opt.isTrain or opt.continue_train:
+            netG = util.load_network(netG, "G", opt.which_epoch, opt)
+            if opt.isTrain:
+                netD = util.load_network(netD, "D", opt.which_epoch, opt)
+            if opt.use_vae:
+                netE = util.load_network(netE, "E", opt.which_epoch, opt)
+
+        return netG, netD, netE
+
+    # preprocess the input, such as moving the tensors to GPUs and
+    # transforming the label map to one-hot encoding
+    # |data|: dictionary of the input data
+
+    def preprocess_input(self, data):
+        # move to GPU and change data types
+        # data['label'] = data['label'].long()
+
+        if not self.opt.isTrain:
+            if self.use_gpu():
+                data["label"] = data["label"].cuda()
+                data["image"] = data["image"].cuda()
+            return data["label"], data["image"], data["image"]
+
+        ## While testing, the input image is the degraded face
+        if self.use_gpu():
+            data["label"] = data["label"].cuda()
+            data["degraded_image"] = data["degraded_image"].cuda()
+            data["image"] = data["image"].cuda()
+
+        # # create one-hot label map
+        # label_map = data['label']
+        # bs, _, h, w = label_map.size()
+        # nc = self.opt.label_nc + 1 if self.opt.contain_dontcare_label \
+        #     else self.opt.label_nc
+        # input_label = self.FloatTensor(bs, nc, h, w).zero_()
+        # input_semantics = input_label.scatter_(1, label_map, 1.0)
+
+        return data["label"], data["image"], data["degraded_image"]
+
+    def compute_generator_loss(self, input_semantics, degraded_image, real_image):
+        G_losses = {}
+
+        fake_image, KLD_loss = self.generate_fake(
+            input_semantics, degraded_image, real_image, compute_kld_loss=self.opt.use_vae
+        )
+
+        if self.opt.use_vae:
+            G_losses["KLD"] = KLD_loss
+
+        pred_fake, pred_real = self.discriminate(input_semantics, fake_image, real_image)
+
+        G_losses["GAN"] = self.criterionGAN(pred_fake, True, for_discriminator=False)
+
+        if not self.opt.no_ganFeat_loss:
+            num_D = len(pred_fake)
+            GAN_Feat_loss = self.FloatTensor(1).fill_(0)
+            for i in range(num_D):  # for each discriminator
+                # last output is the final prediction, so we exclude it
+                num_intermediate_outputs = len(pred_fake[i]) - 1
+                for j in range(num_intermediate_outputs):  # for each layer output
+                    unweighted_loss = self.criterionFeat(pred_fake[i][j], pred_real[i][j].detach())
+                    GAN_Feat_loss += unweighted_loss * self.opt.lambda_feat / num_D
+            G_losses["GAN_Feat"] = GAN_Feat_loss
+
+        if not self.opt.no_vgg_loss:
+            G_losses["VGG"] = self.criterionVGG(fake_image, real_image) * self.opt.lambda_vgg
+
+        return G_losses, fake_image
+
+    def compute_discriminator_loss(self, input_semantics, degraded_image, real_image):
+        D_losses = {}
+        with torch.no_grad():
+            fake_image, _ = self.generate_fake(input_semantics, degraded_image, real_image)
+            fake_image = fake_image.detach()
+            fake_image.requires_grad_()
+
+        pred_fake, pred_real = self.discriminate(input_semantics, fake_image, real_image)
+
+        D_losses["D_Fake"] = self.criterionGAN(pred_fake, False, for_discriminator=True)
+        D_losses["D_real"] = self.criterionGAN(pred_real, True, for_discriminator=True)
+
+        return D_losses
+
+    def encode_z(self, real_image):
+        mu, logvar = self.netE(real_image)
+        z = self.reparameterize(mu, logvar)
+        return z, mu, logvar
+
+    def generate_fake(self, input_semantics, degraded_image, real_image, compute_kld_loss=False):
+        z = None
+        KLD_loss = None
+        if self.opt.use_vae:
+            z, mu, logvar = self.encode_z(real_image)
+            if compute_kld_loss:
+                KLD_loss = self.KLDLoss(mu, logvar) * self.opt.lambda_kld
+
+        fake_image = self.netG(input_semantics, degraded_image, z=z)
+
+        assert (
+            not compute_kld_loss
+        ) or self.opt.use_vae, "You cannot compute KLD loss if opt.use_vae == False"
+
+        return fake_image, KLD_loss
+
+    # Given fake and real image, return the prediction of discriminator
+    # for each fake and real image.
+
+    def discriminate(self, input_semantics, fake_image, real_image):
+
+        if self.opt.no_parsing_map:
+            fake_concat = fake_image
+            real_concat = real_image
+        else:
+            fake_concat = torch.cat([input_semantics, fake_image], dim=1)
+            real_concat = torch.cat([input_semantics, real_image], dim=1)
+
+        # In Batch Normalization, the fake and real images are
+        # recommended to be in the same batch to avoid disparate
+        # statistics in fake and real images.
+        # So both fake and real images are fed to D all at once.
+        fake_and_real = torch.cat([fake_concat, real_concat], dim=0)
+
+        discriminator_out = self.netD(fake_and_real)
+
+        pred_fake, pred_real = self.divide_pred(discriminator_out)
+
+        return pred_fake, pred_real
+
+    # Take the prediction of fake and real images from the combined batch
+    def divide_pred(self, pred):
+        # the prediction contains the intermediate outputs of multiscale GAN,
+        # so it's usually a list
+        if type(pred) == list:
+            fake = []
+            real = []
+            for p in pred:
+                fake.append([tensor[: tensor.size(0) // 2] for tensor in p])
+                real.append([tensor[tensor.size(0) // 2 :] for tensor in p])
+        else:
+            fake = pred[: pred.size(0) // 2]
+            real = pred[pred.size(0) // 2 :]
+
+        return fake, real
+
+    def get_edges(self, t):
+        edge = self.ByteTensor(t.size()).zero_()
+        edge[:, :, :, 1:] = edge[:, :, :, 1:] | (t[:, :, :, 1:] != t[:, :, :, :-1])
+        edge[:, :, :, :-1] = edge[:, :, :, :-1] | (t[:, :, :, 1:] != t[:, :, :, :-1])
+        edge[:, :, 1:, :] = edge[:, :, 1:, :] | (t[:, :, 1:, :] != t[:, :, :-1, :])
+        edge[:, :, :-1, :] = edge[:, :, :-1, :] | (t[:, :, 1:, :] != t[:, :, :-1, :])
+        return edge.float()
+
+    def reparameterize(self, mu, logvar):
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return eps.mul(std) + mu
+
+    def use_gpu(self):
+        return len(self.opt.gpu_ids) > 0

Face_Enhancement/options/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.

Face_Enhancement/options/base_options.py

@@ -0,0 +1,294 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import sys
+import argparse
+import os
+from util import util
+import torch
+import models
+import data
+import pickle
+
+
+class BaseOptions:
+    def __init__(self):
+        self.initialized = False
+
+    def initialize(self, parser):
+        # experiment specifics
+        parser.add_argument(
+            "--name",
+            type=str,
+            default="label2coco",
+            help="name of the experiment. It decides where to store samples and models",
+        )
+
+        parser.add_argument(
+            "--gpu_ids", type=str, default="0", help="gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU"
+        )
+        parser.add_argument(
+            "--checkpoints_dir", type=str, default="./checkpoints", help="models are saved here"
+        )
+        parser.add_argument("--model", type=str, default="pix2pix", help="which model to use")
+        parser.add_argument(
+            "--norm_G",
+            type=str,
+            default="spectralinstance",
+            help="instance normalization or batch normalization",
+        )
+        parser.add_argument(
+            "--norm_D",
+            type=str,
+            default="spectralinstance",
+            help="instance normalization or batch normalization",
+        )
+        parser.add_argument(
+            "--norm_E",
+            type=str,
+            default="spectralinstance",
+            help="instance normalization or batch normalization",
+        )
+        parser.add_argument("--phase", type=str, default="train", help="train, val, test, etc")
+
+        # input/output sizes
+        parser.add_argument("--batchSize", type=int, default=1, help="input batch size")
+        parser.add_argument(
+            "--preprocess_mode",
+            type=str,
+            default="scale_width_and_crop",
+            help="scaling and cropping of images at load time.",
+            choices=(
+                "resize_and_crop",
+                "crop",
+                "scale_width",
+                "scale_width_and_crop",
+                "scale_shortside",
+                "scale_shortside_and_crop",
+                "fixed",
+                "none",
+                "resize",
+            ),
+        )
+        parser.add_argument(
+            "--load_size",
+            type=int,
+            default=1024,
+            help="Scale images to this size. The final image will be cropped to --crop_size.",
+        )
+        parser.add_argument(
+            "--crop_size",
+            type=int,
+            default=512,
+            help="Crop to the width of crop_size (after initially scaling the images to load_size.)",
+        )
+        parser.add_argument(
+            "--aspect_ratio",
+            type=float,
+            default=1.0,
+            help="The ratio width/height. The final height of the load image will be crop_size/aspect_ratio",
+        )
+        parser.add_argument(
+            "--label_nc",
+            type=int,
+            default=182,
+            help="# of input label classes without unknown class. If you have unknown class as class label, specify --contain_dopntcare_label.",
+        )
+        parser.add_argument(
+            "--contain_dontcare_label",
+            action="store_true",
+            help="if the label map contains dontcare label (dontcare=255)",
+        )
+        parser.add_argument("--output_nc", type=int, default=3, help="# of output image channels")
+
+        # for setting inputs
+        parser.add_argument("--dataroot", type=str, default="./datasets/cityscapes/")
+        parser.add_argument("--dataset_mode", type=str, default="coco")
+        parser.add_argument(
+            "--serial_batches",
+            action="store_true",
+            help="if true, takes images in order to make batches, otherwise takes them randomly",
+        )
+        parser.add_argument(
+            "--no_flip",
+            action="store_true",
+            help="if specified, do not flip the images for data argumentation",
+        )
+        parser.add_argument("--nThreads", default=0, type=int, help="# threads for loading data")
+        parser.add_argument(
+            "--max_dataset_size",
+            type=int,
+            default=sys.maxsize,
+            help="Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.",
+        )
+        parser.add_argument(
+            "--load_from_opt_file",
+            action="store_true",
+            help="load the options from checkpoints and use that as default",
+        )
+        parser.add_argument(
+            "--cache_filelist_write",
+            action="store_true",
+            help="saves the current filelist into a text file, so that it loads faster",
+        )
+        parser.add_argument(
+            "--cache_filelist_read", action="store_true", help="reads from the file list cache"
+        )
+
+        # for displays
+        parser.add_argument("--display_winsize", type=int, default=400, help="display window size")
+
+        # for generator
+        parser.add_argument(
+            "--netG", type=str, default="spade", help="selects model to use for netG (pix2pixhd | spade)"
+        )
+        parser.add_argument("--ngf", type=int, default=64, help="# of gen filters in first conv layer")
+        parser.add_argument(
+            "--init_type",
+            type=str,
+            default="xavier",
+            help="network initialization [normal|xavier|kaiming|orthogonal]",
+        )
+        parser.add_argument(
+            "--init_variance", type=float, default=0.02, help="variance of the initialization distribution"
+        )
+        parser.add_argument("--z_dim", type=int, default=256, help="dimension of the latent z vector")
+        parser.add_argument(
+            "--no_parsing_map", action="store_true", help="During training, we do not use the parsing map"
+        )
+
+        # for instance-wise features
+        parser.add_argument(
+            "--no_instance", action="store_true", help="if specified, do *not* add instance map as input"
+        )
+        parser.add_argument(
+            "--nef", type=int, default=16, help="# of encoder filters in the first conv layer"
+        )
+        parser.add_argument("--use_vae", action="store_true", help="enable training with an image encoder.")
+        parser.add_argument(
+            "--tensorboard_log", action="store_true", help="use tensorboard to record the resutls"
+        )
+
+        # parser.add_argument('--img_dir',)
+        parser.add_argument(
+            "--old_face_folder", type=str, default="", help="The folder name of input old face"
+        )
+        parser.add_argument(
+            "--old_face_label_folder", type=str, default="", help="The folder name of input old face label"
+        )
+
+        parser.add_argument("--injection_layer", type=str, default="all", help="")
+
+        self.initialized = True
+        return parser
+
+    def gather_options(self):
+        # initialize parser with basic options
+        if not self.initialized:
+            parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+            parser = self.initialize(parser)
+
+        # get the basic options
+        opt, unknown = parser.parse_known_args()
+
+        # modify model-related parser options
+        model_name = opt.model
+        model_option_setter = models.get_option_setter(model_name)
+        parser = model_option_setter(parser, self.isTrain)
+
+        # modify dataset-related parser options
+        # dataset_mode = opt.dataset_mode
+        # dataset_option_setter = data.get_option_setter(dataset_mode)
+        # parser = dataset_option_setter(parser, self.isTrain)
+
+        opt, unknown = parser.parse_known_args()
+
+        # if there is opt_file, load it.
+        # The previous default options will be overwritten
+        if opt.load_from_opt_file:
+            parser = self.update_options_from_file(parser, opt)
+
+        opt = parser.parse_args()
+        self.parser = parser
+        return opt
+
+    def print_options(self, opt):
+        message = ""
+        message += "----------------- Options ---------------\n"
+        for k, v in sorted(vars(opt).items()):
+            comment = ""
+            default = self.parser.get_default(k)
+            if v != default:
+                comment = "\t[default: %s]" % str(default)
+            message += "{:>25}: {:<30}{}\n".format(str(k), str(v), comment)
+        message += "----------------- End -------------------"
+        # print(message)
+
+    def option_file_path(self, opt, makedir=False):
+        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
+        if makedir:
+            util.mkdirs(expr_dir)
+        file_name = os.path.join(expr_dir, "opt")
+        return file_name
+
+    def save_options(self, opt):
+        file_name = self.option_file_path(opt, makedir=True)
+        with open(file_name + ".txt", "wt") as opt_file:
+            for k, v in sorted(vars(opt).items()):
+                comment = ""
+                default = self.parser.get_default(k)
+                if v != default:
+                    comment = "\t[default: %s]" % str(default)
+                opt_file.write("{:>25}: {:<30}{}\n".format(str(k), str(v), comment))
+
+        with open(file_name + ".pkl", "wb") as opt_file:
+            pickle.dump(opt, opt_file)
+
+    def update_options_from_file(self, parser, opt):
+        new_opt = self.load_options(opt)
+        for k, v in sorted(vars(opt).items()):
+            if hasattr(new_opt, k) and v != getattr(new_opt, k):
+                new_val = getattr(new_opt, k)
+                parser.set_defaults(**{k: new_val})
+        return parser
+
+    def load_options(self, opt):
+        file_name = self.option_file_path(opt, makedir=False)
+        new_opt = pickle.load(open(file_name + ".pkl", "rb"))
+        return new_opt
+
+    def parse(self, save=False):
+
+        opt = self.gather_options()
+        opt.isTrain = self.isTrain  # train or test
+        opt.contain_dontcare_label = False
+
+        self.print_options(opt)
+        if opt.isTrain:
+            self.save_options(opt)
+
+        # Set semantic_nc based on the option.
+        # This will be convenient in many places
+        opt.semantic_nc = (
+            opt.label_nc + (1 if opt.contain_dontcare_label else 0) + (0 if opt.no_instance else 1)
+        )
+
+        # set gpu ids
+        str_ids = opt.gpu_ids.split(",")
+        opt.gpu_ids = []
+        for str_id in str_ids:
+            int_id = int(str_id)
+            if int_id >= 0:
+                opt.gpu_ids.append(int_id)
+
+        if len(opt.gpu_ids) > 0:
+            print("The main GPU is ")
+            print(opt.gpu_ids[0])
+            torch.cuda.set_device(opt.gpu_ids[0])
+
+        assert (
+            len(opt.gpu_ids) == 0 or opt.batchSize % len(opt.gpu_ids) == 0
+        ), "Batch size %d is wrong. It must be a multiple of # GPUs %d." % (opt.batchSize, len(opt.gpu_ids))
+
+        self.opt = opt
+        return self.opt

Face_Enhancement/options/test_options.py

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from .base_options import BaseOptions
+
+
+class TestOptions(BaseOptions):
+    def initialize(self, parser):
+        BaseOptions.initialize(self, parser)
+        parser.add_argument("--results_dir", type=str, default="./results/", help="saves results here.")
+        parser.add_argument(
+            "--which_epoch",
+            type=str,
+            default="latest",
+            help="which epoch to load? set to latest to use latest cached model",
+        )
+        parser.add_argument("--how_many", type=int, default=float("inf"), help="how many test images to run")
+
+        parser.set_defaults(
+            preprocess_mode="scale_width_and_crop", crop_size=256, load_size=256, display_winsize=256
+        )
+        parser.set_defaults(serial_batches=True)
+        parser.set_defaults(no_flip=True)
+        parser.set_defaults(phase="test")
+        self.isTrain = False
+        return parser

Face_Enhancement/requirements.txt

@@ -0,0 +1,9 @@
+torch>=1.0.0
+torchvision
+dominate>=2.3.1
+wandb
+dill
+scikit-image
+tensorboardX
+scipy
+opencv-python

Face_Enhancement/test_face.py

@@ -0,0 +1,45 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import os
+from collections import OrderedDict
+
+import data
+from options.test_options import TestOptions
+from models.pix2pix_model import Pix2PixModel
+from util.visualizer import Visualizer
+import torchvision.utils as vutils
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+opt = TestOptions().parse()
+
+dataloader = data.create_dataloader(opt)
+
+model = Pix2PixModel(opt)
+model.eval()
+
+visualizer = Visualizer(opt)
+
+
+single_save_url = os.path.join(opt.checkpoints_dir, opt.name, opt.results_dir, "each_img")
+
+
+if not os.path.exists(single_save_url):
+    os.makedirs(single_save_url)
+
+
+for i, data_i in enumerate(dataloader):
+    if i * opt.batchSize >= opt.how_many:
+        break
+
+    generated = model(data_i, mode="inference")
+
+    img_path = data_i["path"]
+
+    for b in range(generated.shape[0]):
+        img_name = os.path.split(img_path[b])[-1]
+        save_img_url = os.path.join(single_save_url, img_name)
+
+        vutils.save_image((generated[b] + 1) / 2, save_img_url)
+

Face_Enhancement/util/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.

Face_Enhancement/util/iter_counter.py

@@ -0,0 +1,74 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import os
+import time
+import numpy as np
+
+
+# Helper class that keeps track of training iterations
+class IterationCounter:
+    def __init__(self, opt, dataset_size):
+        self.opt = opt
+        self.dataset_size = dataset_size
+
+        self.first_epoch = 1
+        self.total_epochs = opt.niter + opt.niter_decay
+        self.epoch_iter = 0  # iter number within each epoch
+        self.iter_record_path = os.path.join(self.opt.checkpoints_dir, self.opt.name, "iter.txt")
+        if opt.isTrain and opt.continue_train:
+            try:
+                self.first_epoch, self.epoch_iter = np.loadtxt(
+                    self.iter_record_path, delimiter=",", dtype=int
+                )
+                print("Resuming from epoch %d at iteration %d" % (self.first_epoch, self.epoch_iter))
+            except:
+                print(
+                    "Could not load iteration record at %s. Starting from beginning." % self.iter_record_path
+                )
+
+        self.total_steps_so_far = (self.first_epoch - 1) * dataset_size + self.epoch_iter
+
+    # return the iterator of epochs for the training
+    def training_epochs(self):
+        return range(self.first_epoch, self.total_epochs + 1)
+
+    def record_epoch_start(self, epoch):
+        self.epoch_start_time = time.time()
+        self.epoch_iter = 0
+        self.last_iter_time = time.time()
+        self.current_epoch = epoch
+
+    def record_one_iteration(self):
+        current_time = time.time()
+
+        # the last remaining batch is dropped (see data/__init__.py),
+        # so we can assume batch size is always opt.batchSize
+        self.time_per_iter = (current_time - self.last_iter_time) / self.opt.batchSize
+        self.last_iter_time = current_time
+        self.total_steps_so_far += self.opt.batchSize
+        self.epoch_iter += self.opt.batchSize
+
+    def record_epoch_end(self):
+        current_time = time.time()
+        self.time_per_epoch = current_time - self.epoch_start_time
+        print(
+            "End of epoch %d / %d \t Time Taken: %d sec"
+            % (self.current_epoch, self.total_epochs, self.time_per_epoch)
+        )
+        if self.current_epoch % self.opt.save_epoch_freq == 0:
+            np.savetxt(self.iter_record_path, (self.current_epoch + 1, 0), delimiter=",", fmt="%d")
+            print("Saved current iteration count at %s." % self.iter_record_path)
+
+    def record_current_iter(self):
+        np.savetxt(self.iter_record_path, (self.current_epoch, self.epoch_iter), delimiter=",", fmt="%d")
+        print("Saved current iteration count at %s." % self.iter_record_path)
+
+    def needs_saving(self):
+        return (self.total_steps_so_far % self.opt.save_latest_freq) < self.opt.batchSize
+
+    def needs_printing(self):
+        return (self.total_steps_so_far % self.opt.print_freq) < self.opt.batchSize
+
+    def needs_displaying(self):
+        return (self.total_steps_so_far % self.opt.display_freq) < self.opt.batchSize

Face_Enhancement/util/util.py

@@ -0,0 +1,210 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import re
+import importlib
+import torch
+from argparse import Namespace
+import numpy as np
+from PIL import Image
+import os
+import argparse
+import dill as pickle
+
+
+def save_obj(obj, name):
+    with open(name, "wb") as f:
+        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)
+
+
+def load_obj(name):
+    with open(name, "rb") as f:
+        return pickle.load(f)
+
+
+def copyconf(default_opt, **kwargs):
+    conf = argparse.Namespace(**vars(default_opt))
+    for key in kwargs:
+        print(key, kwargs[key])
+        setattr(conf, key, kwargs[key])
+    return conf
+
+
+# Converts a Tensor into a Numpy array
+# |imtype|: the desired type of the converted numpy array
+def tensor2im(image_tensor, imtype=np.uint8, normalize=True, tile=False):
+    if isinstance(image_tensor, list):
+        image_numpy = []
+        for i in range(len(image_tensor)):
+            image_numpy.append(tensor2im(image_tensor[i], imtype, normalize))
+        return image_numpy
+
+    if image_tensor.dim() == 4:
+        # transform each image in the batch
+        images_np = []
+        for b in range(image_tensor.size(0)):
+            one_image = image_tensor[b]
+            one_image_np = tensor2im(one_image)
+            images_np.append(one_image_np.reshape(1, *one_image_np.shape))
+        images_np = np.concatenate(images_np, axis=0)
+
+        return images_np
+
+    if image_tensor.dim() == 2:
+        image_tensor = image_tensor.unsqueeze(0)
+    image_numpy = image_tensor.detach().cpu().float().numpy()
+    if normalize:
+        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
+    else:
+        image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0
+    image_numpy = np.clip(image_numpy, 0, 255)
+    if image_numpy.shape[2] == 1:
+        image_numpy = image_numpy[:, :, 0]
+    return image_numpy.astype(imtype)
+
+
+# Converts a one-hot tensor into a colorful label map
+def tensor2label(label_tensor, n_label, imtype=np.uint8, tile=False):
+    if label_tensor.dim() == 4:
+        # transform each image in the batch
+        images_np = []
+        for b in range(label_tensor.size(0)):
+            one_image = label_tensor[b]
+            one_image_np = tensor2label(one_image, n_label, imtype)
+            images_np.append(one_image_np.reshape(1, *one_image_np.shape))
+        images_np = np.concatenate(images_np, axis=0)
+        # if tile:
+        #     images_tiled = tile_images(images_np)
+        #     return images_tiled
+        # else:
+        #     images_np = images_np[0]
+        #     return images_np
+        return images_np
+
+    if label_tensor.dim() == 1:
+        return np.zeros((64, 64, 3), dtype=np.uint8)
+    if n_label == 0:
+        return tensor2im(label_tensor, imtype)
+    label_tensor = label_tensor.cpu().float()
+    if label_tensor.size()[0] > 1:
+        label_tensor = label_tensor.max(0, keepdim=True)[1]
+    label_tensor = Colorize(n_label)(label_tensor)
+    label_numpy = np.transpose(label_tensor.numpy(), (1, 2, 0))
+    result = label_numpy.astype(imtype)
+    return result
+
+
+def save_image(image_numpy, image_path, create_dir=False):
+    if create_dir:
+        os.makedirs(os.path.dirname(image_path), exist_ok=True)
+    if len(image_numpy.shape) == 2:
+        image_numpy = np.expand_dims(image_numpy, axis=2)
+    if image_numpy.shape[2] == 1:
+        image_numpy = np.repeat(image_numpy, 3, 2)
+    image_pil = Image.fromarray(image_numpy)
+
+    # save to png
+    image_pil.save(image_path.replace(".jpg", ".png"))
+
+
+def mkdirs(paths):
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            mkdir(path)
+    else:
+        mkdir(paths)
+
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    """
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    """
+    return [atoi(c) for c in re.split("(\d+)", text)]
+
+
+def natural_sort(items):
+    items.sort(key=natural_keys)
+
+
+def str2bool(v):
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("Boolean value expected.")
+
+
+def find_class_in_module(target_cls_name, module):
+    target_cls_name = target_cls_name.replace("_", "").lower()
+    clslib = importlib.import_module(module)
+    cls = None
+    for name, clsobj in clslib.__dict__.items():
+        if name.lower() == target_cls_name:
+            cls = clsobj
+
+    if cls is None:
+        print(
+            "In %s, there should be a class whose name matches %s in lowercase without underscore(_)"
+            % (module, target_cls_name)
+        )
+        exit(0)
+
+    return cls
+
+
+def save_network(net, label, epoch, opt):
+    save_filename = "%s_net_%s.pth" % (epoch, label)
+    save_path = os.path.join(opt.checkpoints_dir, opt.name, save_filename)
+    torch.save(net.cpu().state_dict(), save_path)
+    if len(opt.gpu_ids) and torch.cuda.is_available():
+        net.cuda()
+
+
+def load_network(net, label, epoch, opt):
+    save_filename = "%s_net_%s.pth" % (epoch, label)
+    save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+    save_path = os.path.join(save_dir, save_filename)
+    if os.path.exists(save_path):
+        weights = torch.load(save_path)
+        net.load_state_dict(weights)
+    return net
+
+
+###############################################################################
+# Code from
+# https://github.com/ycszen/pytorch-seg/blob/master/transform.py
+# Modified so it complies with the Citscape label map colors
+###############################################################################
+def uint82bin(n, count=8):
+    """returns the binary of integer n, count refers to amount of bits"""
+    return "".join([str((n >> y) & 1) for y in range(count - 1, -1, -1)])
+
+
+class Colorize(object):
+    def __init__(self, n=35):
+        self.cmap = labelcolormap(n)
+        self.cmap = torch.from_numpy(self.cmap[:n])
+
+    def __call__(self, gray_image):
+        size = gray_image.size()
+        color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0)
+
+        for label in range(0, len(self.cmap)):
+            mask = (label == gray_image[0]).cpu()
+            color_image[0][mask] = self.cmap[label][0]
+            color_image[1][mask] = self.cmap[label][1]
+            color_image[2][mask] = self.cmap[label][2]
+
+        return color_image

Face_Enhancement/util/visualizer.py

@@ -0,0 +1,134 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import os
+import ntpath
+import time
+from . import util
+import scipy.misc
+
+try:
+    from StringIO import StringIO  # Python 2.7
+except ImportError:
+    from io import BytesIO  # Python 3.x
+import torchvision.utils as vutils
+from tensorboardX import SummaryWriter
+import torch
+import numpy as np
+
+
+class Visualizer:
+    def __init__(self, opt):
+        self.opt = opt
+        self.tf_log = opt.isTrain and opt.tf_log
+
+        self.tensorboard_log = opt.tensorboard_log
+
+        self.win_size = opt.display_winsize
+        self.name = opt.name
+        if self.tensorboard_log:
+
+            if self.opt.isTrain:
+                self.log_dir = os.path.join(opt.checkpoints_dir, opt.name, "logs")
+                if not os.path.exists(self.log_dir):
+                    os.makedirs(self.log_dir)
+                self.writer = SummaryWriter(log_dir=self.log_dir)
+            else:
+                print("hi :)")
+                self.log_dir = os.path.join(opt.checkpoints_dir, opt.name, opt.results_dir)
+                if not os.path.exists(self.log_dir):
+                    os.makedirs(self.log_dir)
+
+        if opt.isTrain:
+            self.log_name = os.path.join(opt.checkpoints_dir, opt.name, "loss_log.txt")
+            with open(self.log_name, "a") as log_file:
+                now = time.strftime("%c")
+                log_file.write("================ Training Loss (%s) ================\n" % now)
+
+    # |visuals|: dictionary of images to display or save
+    def display_current_results(self, visuals, epoch, step):
+
+        all_tensor = []
+        if self.tensorboard_log:
+
+            for key, tensor in visuals.items():
+                all_tensor.append((tensor.data.cpu() + 1) / 2)
+
+            output = torch.cat(all_tensor, 0)
+            img_grid = vutils.make_grid(output, nrow=self.opt.batchSize, padding=0, normalize=False)
+
+            if self.opt.isTrain:
+                self.writer.add_image("Face_SPADE/training_samples", img_grid, step)
+            else:
+                vutils.save_image(
+                    output,
+                    os.path.join(self.log_dir, str(step) + ".png"),
+                    nrow=self.opt.batchSize,
+                    padding=0,
+                    normalize=False,
+                )
+
+    # errors: dictionary of error labels and values
+    def plot_current_errors(self, errors, step):
+        if self.tf_log:
+            for tag, value in errors.items():
+                value = value.mean().float()
+                summary = self.tf.Summary(value=[self.tf.Summary.Value(tag=tag, simple_value=value)])
+                self.writer.add_summary(summary, step)
+
+        if self.tensorboard_log:
+
+            self.writer.add_scalar("Loss/GAN_Feat", errors["GAN_Feat"].mean().float(), step)
+            self.writer.add_scalar("Loss/VGG", errors["VGG"].mean().float(), step)
+            self.writer.add_scalars(
+                "Loss/GAN",
+                {
+                    "G": errors["GAN"].mean().float(),
+                    "D": (errors["D_Fake"].mean().float() + errors["D_real"].mean().float()) / 2,
+                },
+                step,
+            )
+
+    # errors: same format as |errors| of plotCurrentErrors
+    def print_current_errors(self, epoch, i, errors, t):
+        message = "(epoch: %d, iters: %d, time: %.3f) " % (epoch, i, t)
+        for k, v in errors.items():
+            v = v.mean().float()
+            message += "%s: %.3f " % (k, v)
+
+        print(message)
+        with open(self.log_name, "a") as log_file:
+            log_file.write("%s\n" % message)
+
+    def convert_visuals_to_numpy(self, visuals):
+        for key, t in visuals.items():
+            tile = self.opt.batchSize > 8
+            if "input_label" == key:
+                t = util.tensor2label(t, self.opt.label_nc + 2, tile=tile)  ## B*H*W*C 0-255 numpy
+            else:
+                t = util.tensor2im(t, tile=tile)
+            visuals[key] = t
+        return visuals
+
+    # save image to the disk
+    def save_images(self, webpage, visuals, image_path):
+        visuals = self.convert_visuals_to_numpy(visuals)
+
+        image_dir = webpage.get_image_dir()
+        short_path = ntpath.basename(image_path[0])
+        name = os.path.splitext(short_path)[0]
+
+        webpage.add_header(name)
+        ims = []
+        txts = []
+        links = []
+
+        for label, image_numpy in visuals.items():
+            image_name = os.path.join(label, "%s.png" % (name))
+            save_path = os.path.join(image_dir, image_name)
+            util.save_image(image_numpy, save_path, create_dir=True)
+
+            ims.append(image_name)
+            txts.append(label)
+            links.append(image_name)
+        webpage.add_images(ims, txts, links, width=self.win_size)

GUI.py

@@ -0,0 +1,217 @@
+import numpy as np
+import cv2
+import PySimpleGUI as sg
+import os.path
+import argparse
+import os
+import sys
+import shutil
+from subprocess import call
+
+def modify(image_filename=None, cv2_frame=None):
+
+    def run_cmd(command):
+        try:
+            call(command, shell=True)
+        except KeyboardInterrupt:
+            print("Process interrupted")
+            sys.exit(1)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_folder", type=str,
+                        default= image_filename, help="Test images")
+    parser.add_argument(
+        "--output_folder",
+        type=str,
+        default="./output",
+        help="Restored images, please use the absolute path",
+    )
+    parser.add_argument("--GPU", type=str, default="-1", help="0,1,2")
+    parser.add_argument(
+        "--checkpoint_name", type=str, default="Setting_9_epoch_100", help="choose which checkpoint"
+    )
+    parser.add_argument("--with_scratch",default="--with_scratch" ,action="store_true")
+    opts = parser.parse_args()
+
+    gpu1 = opts.GPU
+
+    # resolve relative paths before changing directory
+    opts.input_folder = os.path.abspath(opts.input_folder)
+    opts.output_folder = os.path.abspath(opts.output_folder)
+    if not os.path.exists(opts.output_folder):
+        os.makedirs(opts.output_folder)
+
+    main_environment = os.getcwd()
+
+    # Stage 1: Overall Quality Improve
+    print("Running Stage 1: Overall restoration")
+    os.chdir("./Global")
+    stage_1_input_dir = opts.input_folder
+    stage_1_output_dir = os.path.join(
+        opts.output_folder, "stage_1_restore_output")
+    if not os.path.exists(stage_1_output_dir):
+        os.makedirs(stage_1_output_dir)
+
+    if not opts.with_scratch:
+        stage_1_command = (
+            "python test.py --test_mode Full --Quality_restore --test_input "
+            + stage_1_input_dir
+            + " --outputs_dir "
+            + stage_1_output_dir
+            + " --gpu_ids "
+            + gpu1
+        )
+        run_cmd(stage_1_command)
+    else:
+
+        mask_dir = os.path.join(stage_1_output_dir, "masks")
+        new_input = os.path.join(mask_dir, "input")
+        new_mask = os.path.join(mask_dir, "mask")
+        stage_1_command_1 = (
+            "python detection.py --test_path "
+            + stage_1_input_dir
+            + " --output_dir "
+            + mask_dir
+            + " --input_size full_size"
+            + " --GPU "
+            + gpu1
+        )
+        stage_1_command_2 = (
+            "python test.py --Scratch_and_Quality_restore --test_input "
+            + new_input
+            + " --test_mask "
+            + new_mask
+            + " --outputs_dir "
+            + stage_1_output_dir
+            + " --gpu_ids "
+            + gpu1
+        )
+        run_cmd(stage_1_command_1)
+        run_cmd(stage_1_command_2)
+
+    # Solve the case when there is no face in the old photo
+    stage_1_results = os.path.join(stage_1_output_dir, "restored_image")
+    stage_4_output_dir = os.path.join(opts.output_folder, "final_output")
+    if not os.path.exists(stage_4_output_dir):
+        os.makedirs(stage_4_output_dir)
+    for x in os.listdir(stage_1_results):
+        img_dir = os.path.join(stage_1_results, x)
+        shutil.copy(img_dir, stage_4_output_dir)
+
+    print("Finish Stage 1 ...")
+    print("\n")
+
+    # Stage 2: Face Detection
+
+    print("Running Stage 2: Face Detection")
+    os.chdir(".././Face_Detection")
+    stage_2_input_dir = os.path.join(stage_1_output_dir, "restored_image")
+    stage_2_output_dir = os.path.join(
+        opts.output_folder, "stage_2_detection_output")
+    if not os.path.exists(stage_2_output_dir):
+        os.makedirs(stage_2_output_dir)
+    stage_2_command = (
+        "python detect_all_dlib.py --url " + stage_2_input_dir +
+        " --save_url " + stage_2_output_dir
+    )
+    run_cmd(stage_2_command)
+    print("Finish Stage 2 ...")
+    print("\n")
+
+    # Stage 3: Face Restore
+    print("Running Stage 3: Face Enhancement")
+    os.chdir(".././Face_Enhancement")
+    stage_3_input_mask = "./"
+    stage_3_input_face = stage_2_output_dir
+    stage_3_output_dir = os.path.join(
+        opts.output_folder, "stage_3_face_output")
+    if not os.path.exists(stage_3_output_dir):
+        os.makedirs(stage_3_output_dir)
+    stage_3_command = (
+        "python test_face.py --old_face_folder "
+        + stage_3_input_face
+        + " --old_face_label_folder "
+        + stage_3_input_mask
+        + " --tensorboard_log --name "
+        + opts.checkpoint_name
+        + " --gpu_ids "
+        + gpu1
+        + " --load_size 256 --label_nc 18 --no_instance --preprocess_mode resize --batchSize 4 --results_dir "
+        + stage_3_output_dir
+        + " --no_parsing_map"
+    )
+    run_cmd(stage_3_command)
+    print("Finish Stage 3 ...")
+    print("\n")
+
+    # Stage 4: Warp back
+    print("Running Stage 4: Blending")
+    os.chdir(".././Face_Detection")
+    stage_4_input_image_dir = os.path.join(
+        stage_1_output_dir, "restored_image")
+    stage_4_input_face_dir = os.path.join(stage_3_output_dir, "each_img")
+    stage_4_output_dir = os.path.join(opts.output_folder, "final_output")
+    if not os.path.exists(stage_4_output_dir):
+        os.makedirs(stage_4_output_dir)
+    stage_4_command = (
+        "python align_warp_back_multiple_dlib.py --origin_url "
+        + stage_4_input_image_dir
+        + " --replace_url "
+        + stage_4_input_face_dir
+        + " --save_url "
+        + stage_4_output_dir
+    )
+    run_cmd(stage_4_command)
+    print("Finish Stage 4 ...")
+    print("\n")
+
+    print("All the processing is done. Please check the results.")
+
+# --------------------------------- The GUI ---------------------------------
+
+# First the window layout...
+
+images_col = [[sg.Text('Input file:'), sg.In(enable_events=True, key='-IN FILE-'), sg.FileBrowse()],
+              [sg.Button('Modify Photo', key='-MPHOTO-'), sg.Button('Exit')],
+              [sg.Image(filename='', key='-IN-'), sg.Image(filename='', key='-OUT-')],]
+# ----- Full layout -----
+layout = [[sg.VSeperator(), sg.Column(images_col)]]
+
+# ----- Make the window -----
+window = sg.Window('Bringing-old-photos-back-to-life', layout, grab_anywhere=True)
+
+# ----- Run the Event Loop -----
+prev_filename = colorized = cap = None
+while True:
+    event, values = window.read()
+    if event in (None, 'Exit'):
+        break
+
+    elif event == '-MPHOTO-':
+        try:
+            n1 = filename.split("/")[-2]
+            n2 = filename.split("/")[-3]
+            n3 = filename.split("/")[-1]
+            filename= str(f"./{n2}/{n1}")
+            modify(filename)
+           
+            global f_image
+            f_image = f'./output/final_output/{n3}'
+            image = cv2.imread(f_image)
+            window['-OUT-'].update(data=cv2.imencode('.png', image)[1].tobytes())
+            
+        except:
+            continue
+
+    elif event == '-IN FILE-':      # A single filename was chosen
+        filename = values['-IN FILE-']
+        if filename != prev_filename:
+            prev_filename = filename
+            try:
+                image = cv2.imread(filename)
+                window['-IN-'].update(data=cv2.imencode('.png', image)[1].tobytes())
+            except:
+                continue
+
+# ----- Exit program -----
+window.close()

Global/data/Create_Bigfile.py

@@ -0,0 +1,63 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import os
+import struct
+from PIL import Image
+
+IMG_EXTENSIONS = [
+    '.jpg', '.JPG', '.jpeg', '.JPEG',
+    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
+]
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset(dir):
+    images = []
+    assert os.path.isdir(dir), '%s is not a valid directory' % dir
+
+    for root, _, fnames in sorted(os.walk(dir)):
+        for fname in fnames:
+            if is_image_file(fname):
+                #print(fname)
+                path = os.path.join(root, fname)
+                images.append(path)
+
+    return images
+
+### Modify these 3 lines in your own environment
+indir="/home/ziyuwan/workspace/data/temp_old"
+target_folders=['VOC','Real_L_old','Real_RGB_old']
+out_dir ="/home/ziyuwan/workspace/data/temp_old"
+###
+
+if os.path.exists(out_dir) is False:
+    os.makedirs(out_dir)
+
+#
+for target_folder in target_folders:
+    curr_indir = os.path.join(indir, target_folder)
+    curr_out_file = os.path.join(os.path.join(out_dir, '%s.bigfile'%(target_folder)))
+    image_lists = make_dataset(curr_indir)
+    image_lists.sort()
+    with open(curr_out_file, 'wb') as wfid:
+        # write total image number
+        wfid.write(struct.pack('i', len(image_lists)))
+        for i, img_path in enumerate(image_lists):
+             # write file name first
+             img_name = os.path.basename(img_path)
+             img_name_bytes = img_name.encode('utf-8')
+             wfid.write(struct.pack('i', len(img_name_bytes)))
+             wfid.write(img_name_bytes)
+    #
+    #             # write image data in
+             with open(img_path, 'rb') as img_fid:
+                 img_bytes = img_fid.read()
+             wfid.write(struct.pack('i', len(img_bytes)))
+             wfid.write(img_bytes)
+
+             if i % 1000 == 0:
+                 print('write %d images done' % i)

Global/data/Load_Bigfile.py

@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import io
+import os
+import struct
+from PIL import Image
+
+class BigFileMemoryLoader(object):
+    def __load_bigfile(self):
+        print('start load bigfile (%0.02f GB) into memory' % (os.path.getsize(self.file_path)/1024/1024/1024))
+        with open(self.file_path, 'rb') as fid:
+            self.img_num = struct.unpack('i', fid.read(4))[0]
+            self.img_names = []
+            self.img_bytes = []
+            print('find total %d images' % self.img_num)
+            for i in range(self.img_num):
+                img_name_len = struct.unpack('i', fid.read(4))[0]
+                img_name = fid.read(img_name_len).decode('utf-8')
+                self.img_names.append(img_name)
+                img_bytes_len = struct.unpack('i', fid.read(4))[0]
+                self.img_bytes.append(fid.read(img_bytes_len))
+                if i % 5000 == 0:
+                    print('load %d images done' % i)
+            print('load all %d images done' % self.img_num)
+
+    def __init__(self, file_path):
+        super(BigFileMemoryLoader, self).__init__()
+        self.file_path = file_path
+        self.__load_bigfile()
+
+    def __getitem__(self, index):
+        try:
+            img = Image.open(io.BytesIO(self.img_bytes[index])).convert('RGB')
+            return self.img_names[index], img
+        except Exception:
+            print('Image read error for index %d: %s' % (index, self.img_names[index]))
+            return self.__getitem__((index+1)%self.img_num)
+
+
+    def __len__(self):
+        return self.img_num