add merge.py

colab_flask.py

@@ -79,13 +79,13 @@ def call_wav2lip(video_path, audio_path, output_path):
     cmd = [
         "python",
         "wav2lip/inference.py",
-        "--checkpoint_path", checkpoint_path,
+        "--checkpoint_path", checkpoint_path,  #
         # "--segmentation_path", "checkpoints/face_segmentation.pth",
         "--face", video_path,
         "--audio", audio_path,
         "--outfile", output_path,
     ]
-
+    print(f'Call subprocess: {cmd}')
     proc = subprocess.Popen(cmd, shell=False)
     proc.communicate()
     duration = datetime.now() - start
@@ -128,67 +128,30 @@ def call_gfpgan(wav2lip_mp4, audio_path, output_mp4):
         "gfpgan/inference_gfpgan.py",
         "-i", origin_frames_folder,
         "-o", folder_path,
-        "-v", str(1.4),
-        "-s", str(2),
+        # "-v", str(1.4),
+        # "-s", str(2),
         "--only_center_face",
         "--bg_upsampler", 'realesrgan'
     ]
     print(cmd)
-    proc = subprocess.Popen(cmd, shell=False)
+    proc = subprocess.Popen(cmd, shell=True)
     proc.communicate()
     duration = datetime.now() - start
     print(f'inference_gfpgan finished in {duration}')
 
-    restoredFramesPath = os.path.join(folder_path, 'restored_imgs')
-    os.makedirs(restoredFramesPath, exist_ok=True)
-    dir_list = os.listdir(restoredFramesPath)
-    dir_list.sort()
-    batch = 0
-    batchSize = 300
-    for i in tqdm(range(0, len(dir_list), batchSize)):
-        img_array = []
-        start, end = i, i + batchSize
-        print("processing ", start, end)
-        for filename in tqdm(dir_list[start:end]):
-            filename = restoredFramesPath + filename;
-            img = cv2.imread(filename)
-            if img is None:
-                continue
-            height, width, layers = img.shape
-            size = (width, height)
-            img_array.append(img)
-
-        fourcc = cv2.VideoWriter_fourcc(*'X264')
-        filename = 'batch_' + str(batch).zfill(4) + '.mp4'
-        out = cv2.VideoWriter(os.path.join(folder_path, filename),
-                              fourcc, 30, size)
-        batch = batch + 1
-
-        for i in range(len(img_array)):
-            out.write(img_array[i])
-        out.release()
-
-    concatTextFilePath = os.path.join(folder_path, "concat.txt")
-
-    concatTextFile = open(concatTextFilePath, "w")
-    for ips in range(batch):
-        concatTextFile.write("file batch_" + str(ips).zfill(4) + ".mp4\n")
-    concatTextFile.close()
-
-    concatedVideoOutputPath = os.path.join(folder_path, "concated_output.mp4")
-    ff = FFmpeg(
-        inputs={concatTextFilePath: None},
-        outputs={concatedVideoOutputPath: '-y -c copy'}
-    )
-    # !ffmpeg -y -f concat -i {concatTextFilePath} -c copy {concatedVideoOutputPath}
-    ff.run()
-
-    ff = FFmpeg(
-        inputs={concatedVideoOutputPath: None, audio_path: None},
-        outputs={output_mp4: '-y -map 0 -map 1:a -c:v libx264 -c:a aac -shortest'}
-    )
-    # !ffmpeg -y -i {concatedVideoOutputPath} -i {inputAudioPath} -map 0 -map 1:a -c:v copy -shortest {finalProcessedOuputVideo}
-    ff.run()
+    start = datetime.now()
+    cmd = [
+        "python",
+        "merge.py",
+        "-j", folder_path,
+        "-a", audio_path,
+        "-o", output_mp4,
+    ]
+    proc = subprocess.Popen(cmd, shell=True)
+    proc.communicate()
+    duration = datetime.now() - start
+    print(f'Merge output in {duration}')
+    print(output_mp4)
 
     # from google.colab import files
     # files.download(finalProcessedOuputVideo)
@@ -197,20 +160,3 @@ def call_gfpgan(wav2lip_mp4, audio_path, output_mp4):
 if __name__ == '__main__':
     run_with_ngrok(app, auth_token=auth_token)
     app.run()
-
-
-def test():
-    # request
-    import requests
-    ngrok_url = f"http://74c0-34-87-172-60.ngrok-free.app"
-    url = f"{ngrok_url}/wav2lip"
-    print(url)
-    video_path = '/Users/taoluo/Downloads/oIy5B4-vHVw.4.6588496370531551262.0.jpg'
-    audio_path = '/Users/taoluo/Downloads/test_audio.mp3'
-    files = {'video': ('video.jpg', open(video_path, 'rb')), 'audio': ('audio.mp3', open(audio_path, 'rb'))}
-    headers = {'ngrok-skip-browser-warning': 'true'}
-    response = requests.post(url, files=files, headers=headers)
-    # Print the response
-    print(response.json())
-    data = response.json()
-    print(ngrok_url + data['url'])

gfpgan/inference_gfpgan.py

@@ -9,6 +9,8 @@ from tqdm import tqdm
 
 from gfpgan import GFPGANer
 
+root_dir = os.path.dirname(os.path.abspath(__file__))
+
 
 def main():
     """Inference demo for GFPGAN (for users).
@@ -45,11 +47,13 @@ def main():
         help='Image extension. Options: auto | jpg | png, auto means using the same extension as inputs. Default: auto')
     args = parser.parse_args()
 
-    args = parser.parse_args()
-
     # ------------------------ input & output ------------------------
     if args.input.endswith('/'):
         args.input = args.input[:-1]
+    assert os.path.isdir(args.input), f'Input folder {args.input} not exist.'
+
+    os.makedirs(args.output, exist_ok=True)
+
     if os.path.isfile(args.input):
         img_list = [args.input]
     else:
@@ -96,9 +100,9 @@ def main():
         raise ValueError(f'Wrong model version {args.version}.')
 
     # determine model paths
-    model_path = os.path.join('experiments/pretrained_models', model_name + '.pth')
+    model_path = os.path.join(root_dir, 'experiments/pretrained_models', model_name + '.pth')
     if not os.path.isfile(model_path):
-        model_path = os.path.join('realesrgan/weights', model_name + '.pth')
+        model_path = os.path.join(root_dir, 'realesrgan/weights', model_name + '.pth')
     if not os.path.isfile(model_path):
         raise ValueError(f'Model {model_name} does not exist.')
 
@@ -122,21 +126,21 @@ def main():
             input_img, has_aligned=args.aligned, only_center_face=args.only_center_face, paste_back=True)
 
         # save faces
-        if(args.save_faces):
-          for idx, (cropped_face, restored_face) in enumerate(zip(cropped_faces, restored_faces)):
-              # save cropped face
-              save_crop_path = os.path.join(args.output, 'cropped_faces', f'{basename}_{idx:02d}.png')
-              imwrite(cropped_face, save_crop_path)
-              # save restored face
-              if args.suffix is not None:
-                  save_face_name = f'{basename}_{idx:02d}_{args.suffix}.png'
-              else:
-                  save_face_name = f'{basename}_{idx:02d}.png'
-              save_restore_path = os.path.join(args.output, 'restored_faces', save_face_name)
-              imwrite(restored_face, save_restore_path)
-              # save comparison image
-              cmp_img = np.concatenate((cropped_face, restored_face), axis=1)
-              imwrite(cmp_img, os.path.join(args.output, 'cmp', f'{basename}_{idx:02d}.png'))
+        if (args.save_faces):
+            for idx, (cropped_face, restored_face) in enumerate(zip(cropped_faces, restored_faces)):
+                # save cropped face
+                save_crop_path = os.path.join(args.output, 'cropped_faces', f'{basename}_{idx:02d}.png')
+                imwrite(cropped_face, save_crop_path)
+                # save restored face
+                if args.suffix is not None:
+                    save_face_name = f'{basename}_{idx:02d}_{args.suffix}.png'
+                else:
+                    save_face_name = f'{basename}_{idx:02d}.png'
+                save_restore_path = os.path.join(args.output, 'restored_faces', save_face_name)
+                imwrite(restored_face, save_restore_path)
+                # save comparison image
+                cmp_img = np.concatenate((cropped_face, restored_face), axis=1)
+                imwrite(cmp_img, os.path.join(args.output, 'cmp', f'{basename}_{idx:02d}.png'))
 
         # save restored img
         if restored_img is not None:

merge.py

@@ -0,0 +1,77 @@
+import argparse
+import os
+import sys
+
+from tqdm import tqdm
+from ffmpy import FFmpeg
+import cv2
+
+
+def cli(args):
+    folder_path = args.job_folder
+
+    restoredFramesPath = os.path.join(folder_path, 'restored_imgs')
+    os.makedirs(restoredFramesPath, exist_ok=True)
+    dir_list = os.listdir(restoredFramesPath)
+    dir_list.sort()
+    batch = 0
+    batchSize = 300
+    for i in tqdm(range(0, len(dir_list), batchSize)):
+        img_array = []
+        start, end = i, i + batchSize
+        print("processing ", start, end)
+
+        batch_video = os.path.join(folder_path, 'batch_' + str(batch).zfill(4) + '.avi')
+        if os.path.isfile(batch_video):
+            print(f'{batch_video} existed, skipped.')
+            batch = batch + 1
+            continue
+
+        for filename in tqdm(dir_list[start:end]):
+            filename = os.path.join(restoredFramesPath, filename)
+            img = cv2.imread(filename)
+            if img is None:
+                continue
+            height, width, layers = img.shape
+            size = (width, height)
+            img_array.append(img)
+
+        out = cv2.VideoWriter(batch_video, cv2.VideoWriter_fourcc(*'DIVX'), 30, size)
+        batch = batch + 1
+        print(f'batch video dumped: {batch_video}')
+
+        for j in range(len(img_array)):
+            out.write(img_array[j])
+        out.release()
+
+    concatTextFilePath = os.path.join(folder_path, "concat.txt")
+
+    concatTextFile = open(concatTextFilePath, "w")
+    for ips in range(batch):
+        filename = os.path.join(folder_path, "batch_" + str(ips).zfill(4) + ".avi")
+        concatTextFile.write(f"file '{filename}'\n")
+    concatTextFile.close()
+
+    ff = FFmpeg(
+        inputs={
+            concatTextFilePath: '-f concat -safe 0 ',
+            args.audio: None
+        },
+        outputs={args.output: '-y -c:v libx264 -c:a aac -shortest'}
+    )
+    # !ffmpeg -y -f concat -safe 0 -i {concatTextFilePath} -i {audio} -c:v libx264 -c:a aac -shortest {output}
+    ff.run()
+
+
+if __name__ == '__main__':
+    # merge images in restored_imgs folder as video, merge audio
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '-j',
+        '--job_folder',
+        help='job_folder, ./restored_imgs must under it')
+    parser.add_argument('-a', '--audio', type=str, help='audio file path')
+    parser.add_argument('-o', '--output', type=str, help='output file path')
+    args = parser.parse_args()
+    
+    cli(args)