feat(scale_video): keep ratio of video frame

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

tests/test_ffmpeg_tools.py

@@ -13,33 +13,42 @@ 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)
         self.assertEqual(probe['height'], 720)
         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):