diff --git a/gnes/preprocessor/io_utils/video.py b/gnes/preprocessor/io_utils/video.py index 87e4b116..6f7f2998 100644 --- a/gnes/preprocessor/io_utils/video.py +++ b/gnes/preprocessor/io_utils/video.py @@ -18,7 +18,7 @@ from typing import List -from .ffmpeg import parse_media_details, compile_args +from .ffmpeg import parse_media_details, get_media_meta, compile_args from .helper import run_command, run_command_async @@ -147,6 +147,32 @@ def capture_frames(input_fn: str = 'pipe:', raise ValueError( "the buffered video data for stdin should not be empty") + video_meta = get_media_meta(input_fn=input_fn, input_data=input_data) + width = video_meta['frame_width'] + height = video_meta['frame_height'] + + if scale is not None: + _width, _height = map(int, scale.split(':')) + if _width * _height < 0: + if _width > 0: + ratio = _width / width + height = int(ratio * height) + if _height == -2: + height += height % 2 + width = _width + else: + ratio = _height / height + width = int(ratio * width) + if _width == -2: + width += width % 2 + + height = _height + + scale = '%d:%d' % (width, height) + else: + width = _width + height = _height + input_kwargs = { 'err_detect': 'aggressive', 'fflags': 'discardcorrupt' # discard corrupted frames @@ -179,13 +205,6 @@ def capture_frames(input_fn: str = 'pipe:', out, err = run_command( cmd_args, input=input_data, pipe_stdout=True, pipe_stderr=True) - if scale: - width, height = map(int, scale.split(':')) - else: - meta_info = parse_media_details(err.decode()) - width = meta_info['frame_width'] - height = meta_info['frame_height'] - depth = 3 if pix_fmt == 'rgba': depth = 4 diff --git a/tests/test_ffmpeg_tools.py b/tests/test_ffmpeg_tools.py index 668d20b7..5a78fe3c 100644 --- a/tests/test_ffmpeg_tools.py +++ b/tests/test_ffmpeg_tools.py @@ -13,7 +13,7 @@ def setUp(self): self.dirname = os.path.dirname(__file__) self.video_path = os.path.join(self.dirname, 'videos', 'test.mp4') - self.frames = video.capture_frames(input_fn=self.video_path, fps=10, scale='768:360') + self.frames = video.capture_frames(input_fn=self.video_path, fps=10, scale='640:360') def test_probe(self): probe = ffmpeg.probe(self.video_path) @@ -21,25 +21,34 @@ def test_probe(self): self.assertEqual(probe['width'], 1280) self.assertEqual(probe['fps'], 25.0) + + def test_get_media_meta(self): + meta1 = ffmpeg.get_media_meta(input_fn=self.video_path) + with open(self.video_path, 'rb') as f: + data = f.read() + meta2 = ffmpeg.get_media_meta(input_data=data) + self.assertEqual(meta1['frame_width'], meta2['frame_width']) + self.assertEqual(meta1['frame_height'], meta2['frame_height']) + def test_capture_frames(self): - frames1 = video.capture_frames(input_fn=self.video_path, fps=10, scale='768:360') + frames1 = video.capture_frames(input_fn=self.video_path, fps=10, scale='640:-2') with open(self.video_path, 'rb') as f: data = f.read() - frames2 = video.capture_frames(input_data=data, fps=10, scale='768:360') + frames2 = video.capture_frames(input_data=data, fps=10, scale='-1:360') self.assertEqual(frames1.shape, frames2.shape) def test_scale_video(self): - out = video.scale_video(input_fn=self.video_path, scale='768:360') + out = video.scale_video(input_fn=self.video_path, scale='640:360') meta = ffmpeg.get_media_meta(input_data=out, input_options={'format': 'mp4'}) - self.assertEqual(meta['frame_width'], 768) + self.assertEqual(meta['frame_width'], 640) self.assertEqual(meta['frame_height'], 360) def test_encode_video(self): video_data = video.encode_video(images=self.frames) meta = ffmpeg.get_media_meta(input_data=video_data, input_options={'format': 'mp4'}) - self.assertEqual(meta['frame_width'], 768) + self.assertEqual(meta['frame_width'], 640) self.assertEqual(meta['frame_height'], 360) def test_gif_encode(self):