Add M-chip to accelerate adaptation and support for higher version

Face_Detection/align_warp_back_multiple_dlib.py

@@ -111,8 +111,8 @@ def match_histograms(src_image, ref_image):
 
 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
+            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))
@@ -125,7 +125,6 @@ def _origin_face_pts():
 
 
 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
 
@@ -146,7 +145,6 @@ def compute_transformation_matrix(img, landmark, normalize, target_face_scale=1.
 
 
 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
 
@@ -196,8 +194,7 @@ def affine2theta(affine, input_w, input_h, target_w, target_h):
 
 
 def blur_blending(im1, im2, mask):
-
-    mask *= 255.0
+    mask = mask * 255.0
 
     kernel = np.ones((10, 10), np.uint8)
     mask = cv2.erode(mask, kernel, iterations=1)
@@ -215,8 +212,8 @@ def blur_blending(im1, im2, mask):
 
 
 def blur_blending_cv2(im1, im2, mask):
-
-    mask *= 255.0
+    # mask = mask * 255.0
+    mask = mask * 255.0
 
     kernel = np.ones((9, 9), np.uint8)
     mask = cv2.erode(mask, kernel, iterations=3)
@@ -237,14 +234,13 @@ def blur_blending_cv2(im1, im2, mask):
 
 #     Image.composite(
 def Poisson_blending(im1, im2, mask):
-
     # mask=1-mask
-    mask *= 255
+    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 * 255
 
     mask = mask[:, :, 0]
     width, height, channels = im1.shape
@@ -257,8 +253,7 @@ def Poisson_blending(im1, im2, mask):
 
 
 def Poisson_B(im1, im2, mask, center):
-
-    mask *= 255
+    mask = mask * 255
 
     result = cv2.seamlessClone(
         im2.astype("uint8"), im1.astype("uint8"), mask.astype("uint8"), center, cv2.NORMAL_CLONE
@@ -268,7 +263,6 @@ def Poisson_B(im1, im2, mask, center):
 
 
 def seamless_clone(old_face, new_face, raw_mask):
-
     height, width, _ = old_face.shape
     height = height // 2
     width = width // 2
@@ -308,14 +302,16 @@ def seamless_clone(old_face, new_face, raw_mask):
 
 def get_landmark(face_landmarks, id):
     part = face_landmarks.part(id)
-    x = part.x
+    if hasattr(part, 'img'):
+        x = part.img
+    else:
+        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)
@@ -396,7 +392,6 @@ def search(face_landmarks):
             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")
@@ -434,4 +429,3 @@ def search(face_landmarks):
 
         if count % 1000 == 0:
             print("%d have finished ..." % (count))
-

Face_Detection/align_warp_back_multiple_dlib_HR.py

@@ -197,7 +197,7 @@ def affine2theta(affine, input_w, input_h, target_w, target_h):
 
 def blur_blending(im1, im2, mask):
 
-    mask *= 255.0
+    mask = mask * 255.0
 
     kernel = np.ones((10, 10), np.uint8)
     mask = cv2.erode(mask, kernel, iterations=1)
@@ -216,7 +216,7 @@ def blur_blending(im1, im2, mask):
 
 def blur_blending_cv2(im1, im2, mask):
 
-    mask *= 255.0
+    mask = mask * 255.0
 
     kernel = np.ones((9, 9), np.uint8)
     mask = cv2.erode(mask, kernel, iterations=3)
@@ -239,12 +239,13 @@ def blur_blending_cv2(im1, im2, mask):
 def Poisson_blending(im1, im2, mask):
 
     # mask=1-mask
-    mask *= 255
+    mask = mask * 255.0
+
     kernel = np.ones((10, 10), np.uint8)
     mask = cv2.erode(mask, kernel, iterations=1)
     mask /= 255
     mask = 1 - mask
-    mask *= 255
+    mask = mask * 255
 
     mask = mask[:, :, 0]
     width, height, channels = im1.shape
@@ -258,7 +259,7 @@ def Poisson_blending(im1, im2, mask):
 
 def Poisson_B(im1, im2, mask, center):
 
-    mask *= 255
+    mask = mask * 255
 
     result = cv2.seamlessClone(
         im2.astype("uint8"), im1.astype("uint8"), mask.astype("uint8"), center, cv2.NORMAL_CLONE
@@ -308,7 +309,10 @@ def seamless_clone(old_face, new_face, raw_mask):
 
 def get_landmark(face_landmarks, id):
     part = face_landmarks.part(id)
-    x = part.x
+    if hasattr(part, 'img'):
+        x = part.img
+    else:
+        x = part.x
     y = part.y
 
     return (x, y)

Face_Detection/detect_all_dlib.py

@@ -41,7 +41,10 @@ def _origin_face_pts():
 
 def get_landmark(face_landmarks, id):
     part = face_landmarks.part(id)
-    x = part.x
+    if hasattr(part, 'img'):
+        x = part.img
+    else:
+        x = part.x
     y = part.y
 
     return (x, y)

Face_Detection/detect_all_dlib_HR.py

@@ -41,7 +41,10 @@ def _origin_face_pts():
 
 def get_landmark(face_landmarks, id):
     part = face_landmarks.part(id)
-    x = part.x
+    if hasattr(part, 'img'):
+        x = part.img
+    else:
+        x = part.x
     y = part.y
 
     return (x, y)

Face_Enhancement/models/networks/__init__.py

@@ -35,9 +35,12 @@ def modify_commandline_options(parser, is_train):
 def create_network(cls, opt):
     net = cls(opt)
     net.print_network()
-    if len(opt.gpu_ids) > 0:
-        assert torch.cuda.is_available()
-        net.cuda()
+    if opt.mps:
+        net.to("mps")
+    else:
+        if len(opt.gpu_ids) > 0:
+            assert torch.cuda.is_available()
+            net.cuda()
     net.init_weights(opt.init_type, opt.init_variance)
     return net
 

Face_Enhancement/models/pix2pix_model.py

@@ -17,7 +17,6 @@ def __init__(self, opt):
         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
@@ -103,13 +102,20 @@ def preprocess_input(self, data):
         # data['label'] = data['label'].long()
 
         if not self.opt.isTrain:
-            if self.use_gpu():
+            if self.use_mps():
+                data["label"] = data["label"].to(device="mps")
+                data["image"] = data["image"].to(device="mps")
+            elif 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():
+        if self.use_mps():
+            data["label"] = data["label"].to(device="mps")
+            data["degraded_image"] = data["degraded_image"].to(device="mps")
+            data["image"] = data["image"].to(device="mps")
+        elif self.use_gpu():
             data["label"] = data["label"].cuda()
             data["degraded_image"] = data["degraded_image"].cuda()
             data["image"] = data["image"].cuda()
@@ -244,3 +250,5 @@ def reparameterize(self, mu, logvar):
 
     def use_gpu(self):
         return len(self.opt.gpu_ids) > 0
+    def use_mps(self):
+        return self.opt.mps

Face_Enhancement/options/base_options.py

@@ -4,6 +4,8 @@
 import sys
 import argparse
 import os
+
+from numpy.lib._iotools import str2bool
 from util import util
 import torch
 import models
@@ -27,6 +29,9 @@ def initialize(self, parser):
         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(
+            "--mps", type=str2bool, default=False, help="if use M chip, set true"
+        )
         parser.add_argument(
             "--checkpoints_dir", type=str, default="./checkpoints", help="models are saved here"
         )
@@ -281,14 +286,16 @@ def parse(self, save=False):
             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))
+        if opt.mps:
+            torch.device("mps")
+        else:
+            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

Global/detection.py

@@ -13,6 +13,7 @@
 import torch.nn.functional as F
 import torchvision as tv
 from PIL import Image, ImageFile
+from numpy.lib._iotools import str2bool
 
 from detection_models import networks
 from detection_util.util import *
@@ -93,10 +94,13 @@ def main(config):
     model.load_state_dict(checkpoint["model_state"])
     print("model weights loaded")
 
-    if config.GPU >= 0:
-        model.to(config.GPU)
-    else: 
-        model.cpu()
+    if config.mps:
+        model.to("mps")
+    else:
+        if config.GPU >= 0:
+            model.to(config.GPU)
+        else:
+            model.cpu()
     model.eval()
 
     ## dataloader and transformation
@@ -139,11 +143,13 @@ def main(config):
         scratch_image = torch.unsqueeze(scratch_image, 0)
         _, _, ow, oh = scratch_image.shape
         scratch_image_scale = scale_tensor(scratch_image)
-
-        if config.GPU >= 0:
-            scratch_image_scale = scratch_image_scale.to(config.GPU)
+        if config.mps:
+            scratch_image_scale = scratch_image_scale.to("mps")
         else:
-            scratch_image_scale = scratch_image_scale.cpu()
+            if config.GPU >= 0:
+                scratch_image_scale = scratch_image_scale.to(config.GPU)
+            else:
+                scratch_image_scale = scratch_image_scale.cpu()
         with torch.no_grad():
             P = torch.sigmoid(model(scratch_image_scale))
 
@@ -170,6 +176,7 @@ def main(config):
     # parser.add_argument('--checkpoint_name', type=str, default="FT_Epoch_latest.pt", help='Checkpoint Name')
 
     parser.add_argument("--GPU", type=int, default=0)
+    parser.add_argument("--mps", type=str2bool, default=False, help="if use mps acceleration, set true")
     parser.add_argument("--test_path", type=str, default=".")
     parser.add_argument("--output_dir", type=str, default=".")
     parser.add_argument("--input_size", type=str, default="scale_256", help="resize_256|full_size|scale_256")

Global/models/mapping_model.py

@@ -143,11 +143,17 @@ def initialize(self, opt):
             self.netG_A.eval()
             self.netG_B.eval()
 
-        if opt.gpu_ids:
-            self.netG_A.cuda(opt.gpu_ids[0])
-            self.netG_B.cuda(opt.gpu_ids[0])
-            self.mapping_net.cuda(opt.gpu_ids[0])
-
+        if opt.mps:
+            self.netG_A.to('mps')
+            self.netG_B.to('mps')
+            self.mapping_net.to('mps')
+        else:
+            if len(opt.gpu_ids) > 0:
+                # use cuda or use cpu
+                self.netG_A.cuda(opt.gpu_ids[0])
+                self.netG_B.cuda(opt.gpu_ids[0])
+                self.mapping_net.cuda(opt.gpu_ids[0])
+
         if not self.isTrain:
             self.load_network(self.mapping_net, "mapping_net", opt.which_epoch)
 
@@ -324,13 +330,19 @@ def forward(self, label, inst, image, feat, pair=True, infer=False, last_label=N
 
     def inference(self, label, inst):
 
-        use_gpu = len(self.opt.gpu_ids) > 0
-        if use_gpu:
-            input_concat = label.data.cuda()
-            inst_data = inst.cuda()
-        else:
+        if self.opt.mps:
+            label = label.to("mps")
+            inst = inst.to("mps")
             input_concat = label.data
             inst_data = inst
+        else:
+            use_gpu = len(self.opt.gpu_ids) > 0
+            if use_gpu:
+                input_concat = label.data.cuda()
+                inst_data = inst.cuda()
+            else:
+                input_concat = label.data
+                inst_data = inst
 
         label_feat = self.netG_A.forward(input_concat, flow="enc")