| ''' |
| 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 = [] |
| self.audio = None |
| try: |
| 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')) |
| except FileNotFoundError: |
| print('Please install ffmpeg (not ffmpeg-python) by running\n', |
| '\t$ conda install -c conda-forge ffmpeg') |
| sys.exit(0) |
|
|
| 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) |
| 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() |
