From fa50a9d1cf565f5aba5cd66df4c81226cb087e36 Mon Sep 17 00:00:00 2001
From: SWHL <liekkaskono@163.com>
Date: Sat, 8 Jul 2023 19:36:27 +0800
Subject: [PATCH] keep the vsf results under the output directory. Fixed #34

---
 README.md             |  3 ++-
 docs/README_en.md     |  1 +
 rapid_videocr/main.py | 41 ++++++++++++++++++++---------------------
 3 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/README.md b/README.md
index 00d1310..24a2346 100644
--- a/README.md
+++ b/README.md
@@ -56,9 +56,10 @@ flowchart LR
 - [RapidVideOCR高级教程（有python基础的小伙伴）](https://github.com/SWHL/RapidVideOCR/wiki/RapidVideOCR%E9%AB%98%E7%BA%A7%E6%95%99%E7%A8%8B%EF%BC%88%E6%9C%89python%E5%9F%BA%E7%A1%80%E7%9A%84%E5%B0%8F%E4%BC%99%E4%BC%B4%EF%BC%89)
 
 ### 更新日志（[more](https://github.com/SWHL/RapidVideOCR/wiki/%E6%9B%B4%E6%96%B0%E6%97%A5%E5%BF%97)）
-- 🤓2023-07-08 v2.2.2 update:
+- 🤓2023-07-08 v2.2.3 update:
   - 修复批量识别时，不能读取中文路径的问题
   - 修复漏轴时，SRT中跳过问题。目前当出现某一轴未能识别，则会空出位置，便于校对。
+  - 保留VSF识别的中间结果
 - 🐲2023-06-22 Desktop v0.0.3 update:
   - 整合VideoSubFinder界面，增加视频批处理
   - 优化多次选取之后，路径保存问题
diff --git a/docs/README_en.md b/docs/README_en.md
index 55390fe..b79bac5 100644
--- a/docs/README_en.md
+++ b/docs/README_en.md
@@ -56,6 +56,7 @@ flowchart LR
 - 🤓2023-07-08 v2.2.2 update:
    - Fixed the problem that the Chinese path could not be read during batch recognition
    - Skip issue in SRT when fixing missing axes. At present, when a certain axis fails to be recognized, the position will be vacated, which is convenient for proofreading.
+   - Keep the vsf results.
 - 🐲2023-06-22 Desktop v0.0.3 update:
    - Integrate VideoSubFinder interface, increase video batch processing.
    - Optimize the problem of path preservation after multiple selections.
diff --git a/rapid_videocr/main.py b/rapid_videocr/main.py
index 1aa5bc7..15d5679 100644
--- a/rapid_videocr/main.py
+++ b/rapid_videocr/main.py
@@ -2,8 +2,8 @@
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
 import argparse
-import tempfile
 from pathlib import Path
+from typing import Optional
 
 from .rapid_videocr import RapidVideOCR
 from .utils import get_logger
@@ -11,7 +11,7 @@
 
 
 class RapidVideoSubFinderOCR:
-    def __init__(self, vsf_exe_path: str = None, **ocr_params) -> None:
+    def __init__(self, vsf_exe_path: Optional[str] = None, **ocr_params) -> None:
         if vsf_exe_path is None:
             raise ValueError("vsf_exe_path must not be None.")
 
@@ -37,25 +37,24 @@ def __call__(self, video_path: str, output_dir: str = "outputs"):
             self.logger.info(
                 f"[{i+1}/{video_num}] Starting to extract {one_video} key frame"
             )
-            with tempfile.TemporaryDirectory() as tmp_dir:
-                try:
-                    self.vsf(str(one_video), tmp_dir)
-                except Exception as e:
-                    self.logger.error(f"Extract {one_video} error, {e}, skip")
-                    continue
-
-                self.logger.info(f"[{i+1}/{video_num}] Starting to run {one_video} ocr")
-
-                rgb_dir = Path(tmp_dir) / "RGBImages"
-                if not list(rgb_dir.iterdir()):
-                    self.logger.warning(
-                        f"Extracting frames from {one_video} is 0, skip"
-                    )
-                    continue
-
-                save_name = Path(one_video).stem
-                save_srt_dir = Path(output_dir) / save_name
-                self.video_ocr(rgb_dir, save_srt_dir, save_name=save_name)
+
+            save_name = Path(one_video).stem
+            save_dir = Path(output_dir) / save_name
+            save_vsf_dir = save_dir / "VSF_Results"
+
+            try:
+                self.vsf(str(one_video), str(save_vsf_dir))
+            except Exception as e:
+                self.logger.error(f"Extract {one_video} error, {e}, skip")
+                continue
+
+            self.logger.info(f"[{i+1}/{video_num}] Starting to run {one_video} ocr")
+
+            rgb_dir = Path(save_vsf_dir) / "RGBImages"
+            if not list(rgb_dir.iterdir()):
+                self.logger.warning(f"Extracting frames from {one_video} is 0, skip")
+                continue
+            self.video_ocr(rgb_dir, save_dir, save_name=save_name)
 
 
 def main() -> None: