Fix sort errors, and update to v2.1.0

SWHL · Mar 10, 2023 · 073bbb5 · 073bbb5
1 parent 5909f97
commit 073bbb5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/rapid_videocr/rapid_videocr.py b/rapid_videocr/rapid_videocr.py
@@ -54,6 +54,7 @@ def __call__(self,
         mkdir(save_dir)
 
         img_list = list(Path(video_sub_finder_dir).glob('*.jpeg'))
+        img_list = sorted(img_list, key=lambda x: self.get_sort_key(x))
         if not img_list:
             raise RapidVideOCRError(
                 f'{video_sub_finder_dir} has not images with jpeg as suffix.')
@@ -67,6 +68,10 @@ def __call__(self,
 
         self.export_file(save_dir, srt_result, txt_result, out_format)
 
+    @staticmethod
+    def get_sort_key(x):
+        return int(''.join(str(x.stem).split('_')[:4]))
+
     def single_rec(self, img_list: List[str],
                    is_txt_dir: bool) -> Tuple[List, List]:
         srt_result, txt_result = [], []

diff --git a/tests/test_rapid_videocr.py b/tests/test_rapid_videocr.py
@@ -31,9 +31,6 @@ def test_single_rec():
     srt_data = read_txt(srt_path)
     txt_data = read_txt(txt_path)
 
-    print(srt_data)
-    print(txt_data)
-
     assert len(srt_data) == 12
     assert srt_data[2] == '空间里面他绝对赢不了的'
     assert srt_data[-2] == '你们接着善后'
@@ -53,9 +50,6 @@ def test_concat_rec():
     srt_data = read_txt(srt_path)
     txt_data = read_txt(txt_path)
 
-    print(srt_data)
-    print(txt_data)
-
     assert len(srt_data) == 12
     assert srt_data[2] == '空间里面他绝对赢不了的'
     assert srt_data[-2] == '你们接着善后'