|
|
'''
|
|
|
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() |
