Merge pull request #26 from 404Vector/feat/mutli_threading

Feat/multi-threading

tests/_test_i2v2i_.py

@@ -27,7 +27,7 @@ def test_i2v2i(self):
             FrameData(fid, (np.ones((height, width, 3)) * c).astype(np.uint8))
             for fid, c in enumerate(colors)
         ]
-        input_frame_datas.append(FrameData(-1, None))  # end of the frame
+        input_frame_datas.append(None)  # end of the frame
         i2vp = Image2VideoProcessor(
             dst_video_path=config.test_i2v2i["dst_video_path"],
             width=width,
@@ -40,15 +40,15 @@ def test_i2v2i(self):
         frame_datas = input_frame_datas.copy()
         while True:
             frame_data = frame_datas.pop(0)
-            asyncio.run(i2vp(frame_data))
-            if frame_data.frame is None and frame_data.frame_id == -1:
+            i2vp(frame_data)
+            if frame_data is None:
                 break
         v2ip = Video2ImageProcessor(config.test_i2v2i["dst_video_path"])
         while True:
-            output_frame_data = asyncio.run(v2ip())
-            input_frame_data = input_frame_datas[output_frame_data.frame_id]
-            if output_frame_data.frame is None and output_frame_data.frame_id == -1:
+            output_frame_data = v2ip()
+            if output_frame_data is None:
                 break
+            input_frame_data = input_frame_datas[output_frame_data.frame_id]
             self.assertEqual(input_frame_data.frame_id, output_frame_data.frame_id)
             input_avg_color = input_frame_data.frame.mean(axis=0).mean(axis=0)
             output_avg_color = output_frame_data.frame.mean(axis=0).mean(axis=0)
@@ -61,30 +61,3 @@ def test_i2v2i(self):
             f"{i2vp.dst_video_path} is not deleted!",
         )
         return
-        img_input_stream = i2vp.create_stream()
-        try:
-            for fid, input_frame_data in enumerate(input_frame_datas):
-                self.assertEqual(fid, input_frame_data.frame_id)
-                img_input_stream.send(input_frame_data)
-        except StopIteration:
-            pass
-        self.assertEqual(
-            os.path.exists(i2vp.dst_video_path),
-            True,
-            f"{i2vp.dst_video_path} is not exist!",
-        )
-        v2ip = Video2ImageProcessor(config.test_i2v2i["dst_video_path"])
-        img_output_stream = v2ip.create_stream()
-        for fid, output_frame_data in enumerate(img_output_stream):
-            input_frame_data = input_frame_datas[fid]
-            self.assertEqual(input_frame_data.frame_id, output_frame_data.frame_id)
-            input_avg_color = input_frame_data.frame.mean(axis=0).mean(axis=0)
-            output_avg_color = output_frame_data.frame.mean(axis=0).mean(axis=0)
-            davg_color = np.abs(input_avg_color - output_avg_color)
-            self.assertLess(davg_color.mean(), 1)
-        os.remove(i2vp.dst_video_path)
-        self.assertEqual(
-            os.path.exists(i2vp.dst_video_path),
-            False,
-            f"{i2vp.dst_video_path} is not deleted!",
-        )

tests/_test_i2v_.py

@@ -27,7 +27,7 @@ def test_encode_frames(self):
             FrameData(fid, (np.ones((height, width, 3)) * c).astype(np.uint8))
             for fid, c in enumerate(colors)
         ]
-        frame_datas.append(FrameData(-1, None))  # end of the frame
+        frame_datas.append(None)  # end of the frame
         i2vp = Image2VideoProcessor(
             dst_video_path=config.test_i2vp["dst_video_path"],
             width=width,
@@ -38,8 +38,8 @@ def test_encode_frames(self):
         )
         while True:
             frame_data = frame_datas.pop(0)
-            asyncio.run(i2vp(frame_data))
-            if frame_data.frame is None and frame_data.frame_id == -1:
+            i2vp(frame_data)
+            if frame_data is None:
                 break
         self.assertEqual(
             os.path.exists(i2vp.dst_video_path),

tests/_test_v2i2v_.py

@@ -1,9 +1,16 @@
 import asyncio
 import os
+from typing import List
 import unittest
 
 import numpy as np
-from v2v import Video2ImageProcessor, AudioExtractor, Image2VideoProcessor, AudioMerger
+from v2v import (
+    Video2ImageProcessor,
+    AudioExtractor,
+    Image2VideoProcessor,
+    AudioMerger,
+    FrameData,
+)
 from . import _config_ as config
 
 
@@ -23,11 +30,11 @@ def test_v2i2v(self):
             video_path=config.test_v2i2v["test_video_url"],
             ffmpeg_options_output=config.test_v2i2v["v2i_ffmpeg_options_output"],
         )
-        frames = []
+        frames: List[FrameData] = []
         while True:
-            frame = asyncio.run(v2ip())
+            frame = v2ip()
             frames.append(frame)
-            if frame.frame is None and frame.frame_id == -1:
+            if frame is None:
                 break
 
         v2ap = AudioExtractor(
@@ -48,14 +55,14 @@ def test_v2i2v(self):
         try:
             while True:
                 frame_data = frames.pop(0)
-                image = frame_data.frame
-                if frame_data.frame is not None:
+                if frame_data is not None:
+                    image = frame_data.frame
                     image = np.clip(
                         (image.astype(np.int32) - 32) * (128.0 / (128 - 32)), 0, 255
                     )
                     frame_data.frame = image
-                asyncio.run(i2vp(frame_data=frame_data))
-                if frame_data.frame is None and frame_data.frame_id == -1:
+                i2vp(frame_data=frame_data)
+                if frame_data is None:
                     break
         except StopIteration:
             pass

tests/_test_v2i_.py

@@ -23,8 +23,8 @@ def test_extract_frames(self):
 
         progress = 0
         while True:
-            frame_data = asyncio.run(v2ip())
-            if frame_data.frame is None and frame_data.frame_id == -1:
+            frame_data = v2ip()
+            if frame_data is None:
                 break
             self.assertEqual(progress, frame_data.frame_id)
             progress += 1

v2v/_image2video_processor_.py

@@ -39,16 +39,11 @@ def __init__(
         self._progress = 0
         self._is_done = False
 
-    async def __call__(self, frame_data: FrameData):
+    def __call__(self, frame_data: FrameData):
         assert self._is_done is False
-        if frame_data.frame is not None:
+        if frame_data is not None:
             assert self._progress == frame_data.frame_id
-            await asyncio.to_thread(
-                functools.partial(
-                    self._sub_processor.stdin.write,
-                    frame_data.frame.astype(np.uint8).tobytes(),
-                )
-            )
+            self._sub_processor.stdin.write(frame_data.frame.astype(np.uint8).tobytes())
             self._progress += 1
         else:
             self._is_done = True

v2v/_video2image_processor_.py

@@ -24,23 +24,20 @@ def __init__(self, video_path: str, ffmpeg_options_output: Optional[dict] = None
         self._progress = 0
         self._is_done = False
 
-    async def __call__(self) -> FrameData:
+    def __call__(self) -> FrameData:
         assert self._is_done is False
-        frame = await asyncio.to_thread(
-            functools.partial(
-                read_frame_from_process,
-                process=self._sub_processor,
-                width=self._video_info.frame_width,
-                height=self._video_info.frame_height,
-            )
+        frame = read_frame_from_process(
+            process=self._sub_processor,
+            width=self._video_info.frame_width,
+            height=self._video_info.frame_height,
         )
 
         if frame is not None:
             frame_data = FrameData(frame_id=self._progress, frame=frame)
             self._progress += 1
         else:
             self._is_done = True
-            frame_data = FrameData(frame_id=-1, frame=None)
+            frame_data = None
             self._sub_processor.stdout.close()
             self._sub_processor.wait()
         return frame_data