Merge pull request #19 from project-kxkg/jiaen/batch_output

Jiaen/batch output

Former-commit-id: d3f4de05a6facb45767b5eef1fdf6804e04fe29f

pipeline.py

@@ -195,8 +195,10 @@ if os.path.exists(zh_file):
             pass
         count += 1
         segidx = int(count/4)+1
-        
-    with open("{}/{}/{}_en.srt".format(RESULT_PATH, VIDEO_NAME, VIDEO_NAME), "r") as f1, open(temp_file, "a") as f2:
+    en_file = "{}/{}/{}_en.srt".format(RESULT_PATH, VIDEO_NAME, VIDEO_NAME)
+    if args.srt_file is not None:
+        en_file = args.srt_file
+    with open(en_file, "r") as f1, open(temp_file, "a") as f2:
         x = f1.readlines()
         #print(len(x))
         if count >= len(x):
@@ -260,9 +262,10 @@ for sentence, range in tqdm(zip(script_arr, range_arr)):
             flag = True
     # add read-time output back and modify the post-processing by using one batch as an unit.
     srt.set_translation(translate, range, model_name)
-    add_length = srt.check_len_and_split_range(range)
+
+    add_length = srt.check_len_and_split_range(range, threshold)
     srt.realtime_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt",range, add_length,segidx)
-    srt.realtime_bilingual_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_bi.srt",range, add_length,segidx)
+    # srt.realtime_bilingual_write_srt(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_bi.srt",range, add_length,segidx)
 
 # srt.check_len_and_split()
 # srt.write_srt_file_translate(f"{RESULT_PATH}/{VIDEO_NAME}/{VIDEO_NAME}_zh.srt")