update: PIPNet

app.py

@@ -188,19 +188,16 @@ def swap_image(
         if gpu_mode:
             G = G.cuda()
 
-    os.system('ls -hlrt ./')
-    os.system('ls -hlrt ./third_party/PIPNet/FaceBoxesV2/utils/build')
-    os.system('cat ./third_party/PIPNet/FaceBoxesV2/utils/nms/cpu_nms.pyx')
-
+    device = torch.device(0) if gpu_mode else torch.device('cpu')
     source_img = np.array(Image.open(source_image).convert("RGB"))
     net, detector = get_lmk_model()
-    lmk = get_5_from_98(demo_image(source_img, net, detector)[0])
+    lmk = get_5_from_98(demo_image(source_img, net, detector, device=device)[0])
     source_img = norm_crop(source_img, lmk, in_size, mode=align_source, borderValue=0.0)
     source_img = transform(source_img).unsqueeze(0)
 
     target = np.array(Image.open(target_path).convert("RGB"))
     original_target = target.copy()
-    lmk = get_5_from_98(demo_image(target, net, detector)[0])
+    lmk = get_5_from_98(demo_image(target, net, detector, device=device)[0])
     target, M = norm_crop_with_M(target, lmk, in_size, mode=align_target, borderValue=0.0)
     target = transform(target).unsqueeze(0)
     if gpu_mode:
@@ -238,6 +235,7 @@ def process_video(
         G.eval()
         if gpu_mode:
             G = G.cuda()
+    device = torch.device(0) if gpu_mode else torch.device('cpu')
     ''' Target video to frames (.png) '''
     fps = 25.0
     if not os.path.isdir(target_path):
@@ -265,7 +263,7 @@ def process_video(
     target_lmks = []
     for frame_path in tqdm.tqdm(globbed_images):
         target = np.array(Image.open(frame_path).convert("RGB"))
-        lmk = demo_image(target, align_net, align_detector)
+        lmk = demo_image(target, align_net, align_detector, device=device)
         lmk = lmk[0]
         target_lmks.append(lmk)
     ''' Landmark smoothing '''
@@ -283,7 +281,7 @@ def process_video(
     ''' Crop source image '''
     source_img = np.array(Image.open(source_image).convert("RGB"))
     if not use_tddfav2:
-        lmk = get_5_from_98(demo_image(source_img, align_net, align_detector)[0])
+        lmk = get_5_from_98(demo_image(source_img, align_net, align_detector, device=device)[0])
     else:
         lmk = get_lmk(source_img, align_net, align_detector)
     source_img = norm_crop(source_img, lmk, in_size, mode=align_source, borderValue=0.0)