add sound back into enhanced video (#61)

README.md

@@ -81,12 +81,12 @@ python basicsr/setup.py develop
 ### Quick Inference
 
 #### Download Pre-trained Models:
-Download the facelib pretrained models from [[Google Drive](https://drive.google.com/drive/folders/1b_3qwrzY_kTQh0-SnBoGBgOrJ_PLZSKm?usp=sharing) | [OneDrive](https://entuedu-my.sharepoint.com/:f:/g/personal/s200094_e_ntu_edu_sg/EvDxR7FcAbZMp_MA9ouq7aQB8XTppMb3-T0uGZ_2anI2mg?e=DXsJFo)] to the `weights/facelib` folder. You can manually download the pretrained models OR download by runing the following command.
+Download the facelib pretrained models from [[Google Drive](https://drive.google.com/drive/folders/1b_3qwrzY_kTQh0-SnBoGBgOrJ_PLZSKm?usp=sharing) | [OneDrive](https://entuedu-my.sharepoint.com/:f:/g/personal/s200094_e_ntu_edu_sg/EvDxR7FcAbZMp_MA9ouq7aQB8XTppMb3-T0uGZ_2anI2mg?e=DXsJFo)] to the `weights/facelib` folder. You can manually download the pretrained models OR download by running the following command.
 ```
 python scripts/download_pretrained_models.py facelib
 ```
 
-Download the CodeFormer pretrained models from [[Google Drive](https://drive.google.com/drive/folders/1CNNByjHDFt0b95q54yMVp6Ifo5iuU6QS?usp=sharing) | [OneDrive](https://entuedu-my.sharepoint.com/:f:/g/personal/s200094_e_ntu_edu_sg/EoKFj4wo8cdIn2-TY2IV6CYBhZ0pIG4kUOeHdPR_A5nlbg?e=AO8UN9)] to the `weights/CodeFormer` folder. You can manually download the pretrained models OR download by runing the following command.
+Download the CodeFormer pretrained models from [[Google Drive](https://drive.google.com/drive/folders/1CNNByjHDFt0b95q54yMVp6Ifo5iuU6QS?usp=sharing) | [OneDrive](https://entuedu-my.sharepoint.com/:f:/g/personal/s200094_e_ntu_edu_sg/EoKFj4wo8cdIn2-TY2IV6CYBhZ0pIG4kUOeHdPR_A5nlbg?e=AO8UN9)] to the `weights/CodeFormer` folder. You can manually download the pretrained models OR download by running the following command.
 ```
 python scripts/download_pretrained_models.py CodeFormer
 ```
@@ -115,8 +115,7 @@ python inference_codeformer.py -w 0.7 --input_path [image folder/image path]
 :clapper: Video Enhancement
 ```
 # For video clips
-# Set frame rate of saved video via '--save_video_fps 24'
-python inference_codeformer.py --bg_upsampler realesrgan --face_upsample -w 1.0 --input_path [video path] --save_video_fps 24
+python inference_codeformer.py --bg_upsampler realesrgan --face_upsample -w 1.0 --input_path [video path]
 ```
 
 

basicsr/utils/video_util.py

@@ -0,0 +1,119 @@
+'''
+The code is modified from the Real-ESRGAN:
+https://github.com/xinntao/Real-ESRGAN/blob/master/inference_realesrgan_video.py
+
+'''
+import cv2
+import sys
+import numpy as np
+
+try:
+    import ffmpeg
+except ImportError:
+    import pip
+    pip.main(['install', '--user', 'ffmpeg-python'])
+    import ffmpeg
+
+def get_video_meta_info(video_path):
+    ret = {}
+    probe = ffmpeg.probe(video_path)
+    video_streams = [stream for stream in probe['streams'] if stream['codec_type'] == 'video']
+    has_audio = any(stream['codec_type'] == 'audio' for stream in probe['streams'])
+    ret['width'] = video_streams[0]['width']
+    ret['height'] = video_streams[0]['height']
+    ret['fps'] = eval(video_streams[0]['avg_frame_rate'])
+    ret['audio'] = ffmpeg.input(video_path).audio if has_audio else None
+    ret['nb_frames'] = int(video_streams[0]['nb_frames'])
+    return ret
+
+class VideoReader:
+    def __init__(self, video_path):
+        self.paths = []  # for image&folder type
+        self.audio = None
+        self.stream_reader = (
+            ffmpeg.input(video_path).output('pipe:', format='rawvideo', pix_fmt='bgr24',
+                                            loglevel='error').run_async(
+                                                pipe_stdin=True, pipe_stdout=True, cmd='ffmpeg'))
+        meta = get_video_meta_info(video_path)
+        self.width = meta['width']
+        self.height = meta['height']
+        self.input_fps = meta['fps']
+        self.audio = meta['audio']
+        self.nb_frames = meta['nb_frames']
+
+        self.idx = 0
+
+    def get_resolution(self):
+        return self.height, self.width
+
+    def get_fps(self):
+        if self.input_fps is not None:
+            return self.input_fps
+        return 24
+
+    def get_audio(self):
+        return self.audio
+
+    def __len__(self):
+        return self.nb_frames
+
+    def get_frame_from_stream(self):
+        img_bytes = self.stream_reader.stdout.read(self.width * self.height * 3)  # 3 bytes for one pixel
+        if not img_bytes:
+            return None
+        img = np.frombuffer(img_bytes, np.uint8).reshape([self.height, self.width, 3])
+        return img
+
+    def get_frame_from_list(self):
+        if self.idx >= self.nb_frames:
+            return None
+        img = cv2.imread(self.paths[self.idx])
+        self.idx += 1
+        return img
+
+    def get_frame(self):
+        return self.get_frame_from_stream()
+
+
+    def close(self):
+        self.stream_reader.stdin.close()
+        self.stream_reader.wait()
+
+
+class VideoWriter:
+    def __init__(self, video_save_path, height, width, fps, audio):
+        if height > 2160:
+            print('You are generating video that is larger than 4K, which will be very slow due to IO speed.',
+                  'We highly recommend to decrease the outscale(aka, -s).')
+        if audio is not None:
+            self.stream_writer = (
+                ffmpeg.input('pipe:', format='rawvideo', pix_fmt='bgr24', s=f'{width}x{height}',
+                            framerate=fps).output(
+                                audio,
+                                video_save_path,
+                                pix_fmt='yuv420p',
+                                vcodec='libx264',
+                                loglevel='error',
+                                acodec='copy').overwrite_output().run_async(
+                                    pipe_stdin=True, pipe_stdout=True, cmd='ffmpeg'))
+        else:
+            self.stream_writer = (
+                ffmpeg.input('pipe:', format='rawvideo', pix_fmt='bgr24', s=f'{width}x{height}',
+                            framerate=fps).output(
+                                video_save_path, pix_fmt='yuv420p', vcodec='libx264',
+                                loglevel='error').overwrite_output().run_async(
+                                    pipe_stdin=True, pipe_stdout=True, cmd='ffmpeg'))
+
+    def write_frame(self, frame):
+        try:
+            frame = frame.astype(np.uint8).tobytes()
+            self.stream_writer.stdin.write(frame)
+        except BrokenPipeError:
+            print('Please re-install ffmpeg and libx264 by running\n',
+                  '\t$ conda install -c conda-forge ffmpeg\n',
+                  '\t$ conda install -c conda-forge x264')
+            sys.exit(0)
+
+    def close(self):
+        self.stream_writer.stdin.close()
+        self.stream_writer.wait()

facelib/utils/misc.py

@@ -7,13 +7,13 @@ import torch
 from torch.hub import download_url_to_file, get_dir
 from urllib.parse import urlparse
 # from basicsr.utils.download_util import download_file_from_google_drive
-import gdown
-
 
 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 
 def download_pretrained_models(file_ids, save_path_root):
+    import gdown
+    
     os.makedirs(save_path_root, exist_ok=True)
 
     for file_name, file_id in file_ids.items():

inference_codeformer.py

@@ -72,7 +72,7 @@ if __name__ == '__main__':
     parser.add_argument('--face_upsample', action='store_true', help='Face upsampler after enhancement. Default: False')
     parser.add_argument('--bg_tile', type=int, default=400, help='Tile size for background sampler. Default: 400')
     parser.add_argument('--suffix', type=str, default=None, help='Suffix of the restored faces. Default: None')
-    parser.add_argument('--save_video_fps', type=int, default=24, help='Frame rate for saving video. Default: 24')
+    parser.add_argument('--save_video_fps', type=float, default=None, help='Frame rate for saving video. Default: None')
 
     args = parser.parse_args()
 
@@ -83,15 +83,19 @@ if __name__ == '__main__':
         input_img_list = [args.input_path]
         result_root = f'results/test_img_{w}'
     elif args.input_path.endswith(('mp4', 'mov', 'avi')): # input video path
+        from basicsr.utils.video_util import VideoReader, VideoWriter
         input_img_list = []
-        vidcap = cv2.VideoCapture(args.input_path)
-        success, image = vidcap.read()
-        while success:
+        vidreader = VideoReader(args.input_path)
+        image = vidreader.get_frame()
+        while image is not None:
             input_img_list.append(image)
-            success, image = vidcap.read()
-        input_video = True
+            image = vidreader.get_frame()
+        audio = vidreader.get_audio()
+        fps = vidreader.get_fps() if args.save_video_fps is None else args.save_video_fps   
         video_name = os.path.basename(args.input_path)[:-4]
         result_root = f'results/{video_name}_{w}'
+        input_video = True
+        vidreader.close()
     else: # input img folder
         if args.input_path.endswith('/'):  # solve when path ends with /
             args.input_path = args.input_path[:-1]
@@ -241,6 +245,7 @@ if __name__ == '__main__':
 
     # save enhanced video
     if input_video:
+        print('Video Saving...')
         # load images
         video_frames = []
         img_list = sorted(glob.glob(os.path.join(result_root, 'final_results', '*.[jp][pn]g')))
@@ -248,14 +253,14 @@ if __name__ == '__main__':
             img = cv2.imread(img_path)
             video_frames.append(img)
         # write images to video
-        h, w = video_frames[0].shape[:2]
+        height, width = video_frames[0].shape[:2]
         if args.suffix is not None:
             video_name = f'{video_name}_{args.suffix}.png'
         save_restore_path = os.path.join(result_root, f'{video_name}.mp4')
-        writer = cv2.VideoWriter(save_restore_path, cv2.VideoWriter_fourcc(*"mp4v"),
-                                    args.save_video_fps, (w, h))            
+        vidwriter = VideoWriter(save_restore_path, height, width, fps, audio)
+         
         for f in video_frames:
-            writer.write(f)
-        writer.release()
+            vidwriter.write_frame(f)
+        vidwriter.close()
 
     print(f'\nAll results are saved in {result_root}')

requirements.txt

@@ -15,6 +15,4 @@ tqdm
 yapf
 lpips
 gdown # supports downloading the large file from Google Drive
-# cmake
-# dlib
-# conda install -c conda-forge dlib
+ffmpeg-python