Merge pull request #54 from jhj0517/add-raw-type

app.py

@@ -50,7 +50,7 @@ class App:
                                                label="Model")
                         dd_lang = gr.Dropdown(choices=["Automatic Detection"] + self.whisper_inf.available_langs,
                                               value="Automatic Detection", label="Language")
-                        dd_subformat = gr.Dropdown(["SRT", "WebVTT"], value="SRT", label="Subtitle Format")
+                        dd_file_format = gr.Dropdown(["SRT", "WebVTT", "txt"], value="SRT", label="File Format")
                     with gr.Row():
                         cb_translate = gr.Checkbox(value=False, label="Translate to English?", interactive=True)
                     with gr.Row():
@@ -66,7 +66,7 @@ class App:
                         tb_indicator = gr.Textbox(label="Output", scale=8)
                         btn_openfolder = gr.Button('📂', scale=2)
 
-                    params = [input_file, dd_model, dd_lang, dd_subformat, cb_translate, cb_timestamp]
+                    params = [input_file, dd_model, dd_lang, dd_file_format, cb_translate, cb_timestamp]
                     advanced_params = [nb_beam_size, nb_log_prob_threshold, nb_no_speech_threshold, dd_compute_type]
                     btn_run.click(fn=self.whisper_inf.transcribe_file,
                                   inputs=params + advanced_params,
@@ -88,7 +88,7 @@ class App:
                                                label="Model")
                         dd_lang = gr.Dropdown(choices=["Automatic Detection"] + self.whisper_inf.available_langs,
                                               value="Automatic Detection", label="Language")
-                        dd_subformat = gr.Dropdown(choices=["SRT", "WebVTT"], value="SRT", label="Subtitle Format")
+                        dd_file_format = gr.Dropdown(choices=["SRT", "WebVTT", "txt"], value="SRT", label="File Format")
                     with gr.Row():
                         cb_translate = gr.Checkbox(value=False, label="Translate to English?", interactive=True)
                     with gr.Row():
@@ -105,7 +105,7 @@ class App:
                         tb_indicator = gr.Textbox(label="Output", scale=8)
                         btn_openfolder = gr.Button('📂', scale=2)
 
-                    params = [tb_youtubelink, dd_model, dd_lang, dd_subformat, cb_translate, cb_timestamp]
+                    params = [tb_youtubelink, dd_model, dd_lang, dd_file_format, cb_translate, cb_timestamp]
                     advanced_params = [nb_beam_size, nb_log_prob_threshold, nb_no_speech_threshold, dd_compute_type]
                     btn_run.click(fn=self.whisper_inf.transcribe_youtube,
                                   inputs=params + advanced_params,
@@ -123,7 +123,7 @@ class App:
                                                label="Model")
                         dd_lang = gr.Dropdown(choices=["Automatic Detection"] + self.whisper_inf.available_langs,
                                               value="Automatic Detection", label="Language")
-                        dd_subformat = gr.Dropdown(["SRT", "WebVTT"], value="SRT", label="Subtitle Format")
+                        dd_file_format = gr.Dropdown(["SRT", "WebVTT", "txt"], value="SRT", label="File Format")
                     with gr.Row():
                         cb_translate = gr.Checkbox(value=False, label="Translate to English?", interactive=True)
                     with gr.Accordion("Advanced_Parameters", open=False):
@@ -137,7 +137,7 @@ class App:
                         tb_indicator = gr.Textbox(label="Output", scale=8)
                         btn_openfolder = gr.Button('📂', scale=2)
 
-                    params = [mic_input, dd_model, dd_lang, dd_subformat, cb_translate]
+                    params = [mic_input, dd_model, dd_lang, dd_file_format, cb_translate]
                     advanced_params = [nb_beam_size, nb_log_prob_threshold, nb_no_speech_threshold, dd_compute_type]
                     btn_run.click(fn=self.whisper_inf.transcribe_mic,
                                   inputs=params + advanced_params,

modules/faster_whisper_inference.py

@@ -13,7 +13,7 @@ import torch
 import gradio as gr
 
 from .base_interface import BaseInterface
-from modules.subtitle_manager import get_srt, get_vtt, write_file, safe_filename
+from modules.subtitle_manager import get_srt, get_vtt, get_txt, write_file, safe_filename
 from modules.youtube_manager import get_ytdata, get_ytaudio
 
 
@@ -34,7 +34,7 @@ class FasterWhisperInference(BaseInterface):
                         fileobjs: list,
                         model_size: str,
                         lang: str,
-                        subformat: str,
+                        file_format: str,
                         istranslate: bool,
                         add_timestamp: bool,
                         beam_size: int,
@@ -54,8 +54,8 @@ class FasterWhisperInference(BaseInterface):
             Whisper model size from gr.Dropdown()
         lang: str
             Source language of the file to transcribe from gr.Dropdown()
-        subformat: str
-            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        file_format: str
+            File format to write from gr.Dropdown(). Supported format: [SRT, WebVTT, txt]
         istranslate: bool
             Boolean value from gr.Checkbox() that determines whether to translate to English.
             It's Whisper's feature to translate speech from another language directly into English end-to-end.
@@ -97,12 +97,13 @@ class FasterWhisperInference(BaseInterface):
 
                 file_name, file_ext = os.path.splitext(os.path.basename(fileobj.orig_name))
                 file_name = safe_filename(file_name)
-                subtitle = self.generate_and_write_subtitle(
+                subtitle = self.generate_and_write_file(
                     file_name=file_name,
                     transcribed_segments=transcribed_segments,
                     add_timestamp=add_timestamp,
-                    subformat=subformat
+                    file_format=file_format
                 )
+                print(f"{subtitle}")
                 files_info[file_name] = {"subtitle": subtitle, "time_for_task": time_for_task}
 
             total_result = ''
@@ -125,7 +126,7 @@ class FasterWhisperInference(BaseInterface):
                            youtubelink: str,
                            model_size: str,
                            lang: str,
-                           subformat: str,
+                           file_format: str,
                            istranslate: bool,
                            add_timestamp: bool,
                            beam_size: int,
@@ -145,8 +146,8 @@ class FasterWhisperInference(BaseInterface):
             Whisper model size from gr.Dropdown()
         lang: str
             Source language of the file to transcribe from gr.Dropdown()
-        subformat: str
-            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        file_format: str
+            File format to write from gr.Dropdown(). Supported format: [SRT, WebVTT, txt]
         istranslate: bool
             Boolean value from gr.Checkbox() that determines whether to translate to English.
             It's Whisper's feature to translate speech from another language directly into English end-to-end.
@@ -191,11 +192,11 @@ class FasterWhisperInference(BaseInterface):
             progress(1, desc="Completed!")
 
             file_name = safe_filename(yt.title)
-            subtitle = self.generate_and_write_subtitle(
+            subtitle = self.generate_and_write_file(
                 file_name=file_name,
                 transcribed_segments=transcribed_segments,
                 add_timestamp=add_timestamp,
-                subformat=subformat
+                file_format=file_format
             )
             return f"Done in {self.format_time(time_for_task)}! Subtitle file is in the outputs folder.\n\n{subtitle}"
         except Exception as e:
@@ -217,7 +218,7 @@ class FasterWhisperInference(BaseInterface):
                        micaudio: str,
                        model_size: str,
                        lang: str,
-                       subformat: str,
+                       file_format: str,
                        istranslate: bool,
                        beam_size: int,
                        log_prob_threshold: float,
@@ -236,8 +237,8 @@ class FasterWhisperInference(BaseInterface):
             Whisper model size from gr.Dropdown()
         lang: str
             Source language of the file to transcribe from gr.Dropdown()
-        subformat: str
-            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        file_format: str
+            File format to write from gr.Dropdown(). Supported format: [SRT, WebVTT, txt]
         istranslate: bool
             Boolean value from gr.Checkbox() that determines whether to translate to English.
             It's Whisper's feature to translate speech from another language directly into English end-to-end.
@@ -276,11 +277,11 @@ class FasterWhisperInference(BaseInterface):
             )
             progress(1, desc="Completed!")
 
-            subtitle = self.generate_and_write_subtitle(
+            subtitle = self.generate_and_write_file(
                 file_name="Mic",
                 transcribed_segments=transcribed_segments,
                 add_timestamp=True,
-                subformat=subformat
+                file_format=file_format
             )
             return f"Done in {self.format_time(time_for_task)}! Subtitle file is in the outputs folder.\n\n{subtitle}"
         except Exception as e:
@@ -378,11 +379,11 @@ class FasterWhisperInference(BaseInterface):
             )
 
     @staticmethod
-    def generate_and_write_subtitle(file_name: str,
-                                    transcribed_segments: list,
-                                    add_timestamp: bool,
-                                    subformat: str,
-                                    ) -> str:
+    def generate_and_write_file(file_name: str,
+                                transcribed_segments: list,
+                                add_timestamp: bool,
+                                file_format: str,
+                                ) -> str:
         """
         This method writes subtitle file and returns str to gr.Textbox
         """
@@ -392,13 +393,18 @@ class FasterWhisperInference(BaseInterface):
         else:
             output_path = os.path.join("outputs", f"{file_name}")
 
-        if subformat == "SRT":
-            subtitle = get_srt(transcribed_segments)
-            write_file(subtitle, f"{output_path}.srt")
-        elif subformat == "WebVTT":
-            subtitle = get_vtt(transcribed_segments)
-            write_file(subtitle, f"{output_path}.vtt")
-        return subtitle
+        if file_format == "SRT":
+            content = get_srt(transcribed_segments)
+            write_file(content, f"{output_path}.srt")
+
+        elif file_format == "WebVTT":
+            content = get_vtt(transcribed_segments)
+            write_file(content, f"{output_path}.vtt")
+
+        elif file_format == "txt":
+            content = get_txt(transcribed_segments)
+            write_file(content, f"{output_path}.txt")
+        return content
 
     @staticmethod
     def format_time(elapsed_time: float) -> str:

modules/subtitle_manager.py

@@ -44,6 +44,15 @@ def get_vtt(segments):
     return output
 
 
+def get_txt(segments):
+    output = ""
+    for i, segment in enumerate(segments):
+        if segment['text'].startswith(' '):
+            segment['text'] = segment['text'][1:]
+        output += f"{segment['text']}\n"
+    return output
+
+
 def parse_srt(file_path):
     """Reads SRT file and returns as dict"""
     with open(file_path, 'r', encoding='utf-8') as file:

modules/whisper_Inference.py

@@ -8,7 +8,7 @@ from datetime import datetime
 import torch
 
 from .base_interface import BaseInterface
-from modules.subtitle_manager import get_srt, get_vtt, write_file, safe_filename
+from modules.subtitle_manager import get_srt, get_vtt, get_txt, write_file, safe_filename
 from modules.youtube_manager import get_ytdata, get_ytaudio
 
 DEFAULT_MODEL_SIZE = "large-v2"
@@ -30,7 +30,7 @@ class WhisperInference(BaseInterface):
                         fileobjs: list,
                         model_size: str,
                         lang: str,
-                        subformat: str,
+                        file_format: str,
                         istranslate: bool,
                         add_timestamp: bool,
                         beam_size: int,
@@ -49,8 +49,8 @@ class WhisperInference(BaseInterface):
             Whisper model size from gr.Dropdown()
         lang: str
             Source language of the file to transcribe from gr.Dropdown()
-        subformat: str
-            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        file_format: str
+            File format to write from gr.Dropdown(). Supported format: [SRT, WebVTT, txt]
         istranslate: bool
             Boolean value from gr.Checkbox() that determines whether to translate to English.
             It's Whisper's feature to translate speech from another language directly into English end-to-end.
@@ -93,11 +93,11 @@ class WhisperInference(BaseInterface):
 
                 file_name, file_ext = os.path.splitext(os.path.basename(fileobj.orig_name))
                 file_name = safe_filename(file_name)
-                subtitle = self.generate_and_write_subtitle(
+                subtitle = self.generate_and_write_file(
                     file_name=file_name,
                     transcribed_segments=result,
                     add_timestamp=add_timestamp,
-                    subformat=subformat
+                    file_format=file_format
                 )
 
                 files_info[file_name] = {"subtitle": subtitle, "elapsed_time": elapsed_time}
@@ -122,7 +122,7 @@ class WhisperInference(BaseInterface):
                            youtubelink: str,
                            model_size: str,
                            lang: str,
-                           subformat: str,
+                           file_format: str,
                            istranslate: bool,
                            add_timestamp: bool,
                            beam_size: int,
@@ -141,8 +141,8 @@ class WhisperInference(BaseInterface):
             Whisper model size from gr.Dropdown()
         lang: str
             Source language of the file to transcribe from gr.Dropdown()
-        subformat: str
-            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        file_format: str
+            File format to write from gr.Dropdown(). Supported format: [SRT, WebVTT, txt]
         istranslate: bool
             Boolean value from gr.Checkbox() that determines whether to translate to English.
             It's Whisper's feature to translate speech from another language directly into English end-to-end.
@@ -181,11 +181,11 @@ class WhisperInference(BaseInterface):
             progress(1, desc="Completed!")
 
             file_name = safe_filename(yt.title)
-            subtitle = self.generate_and_write_subtitle(
+            subtitle = self.generate_and_write_file(
                 file_name=file_name,
                 transcribed_segments=result,
                 add_timestamp=add_timestamp,
-                subformat=subformat
+                file_format=file_format
             )
 
             return f"Done in {self.format_time(elapsed_time)}! Subtitle file is in the outputs folder.\n\n{subtitle}"
@@ -209,7 +209,7 @@ class WhisperInference(BaseInterface):
                        micaudio: str,
                        model_size: str,
                        lang: str,
-                       subformat: str,
+                       file_format: str,
                        istranslate: bool,
                        beam_size: int,
                        log_prob_threshold: float,
@@ -227,8 +227,8 @@ class WhisperInference(BaseInterface):
             Whisper model size from gr.Dropdown()
         lang: str
             Source language of the file to transcribe from gr.Dropdown()
-        subformat: str
-            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT]
+        file_format: str
+            Subtitle format to write from gr.Dropdown(). Supported format: [SRT, WebVTT, txt]
         istranslate: bool
             Boolean value from gr.Checkbox() that determines whether to translate to English.
             It's Whisper's feature to translate speech from another language directly into English end-to-end.
@@ -261,11 +261,11 @@ class WhisperInference(BaseInterface):
                                                    progress=progress)
             progress(1, desc="Completed!")
 
-            subtitle = self.generate_and_write_subtitle(
+            subtitle = self.generate_and_write_file(
                 file_name="Mic",
                 transcribed_segments=result,
                 add_timestamp=True,
-                subformat=subformat
+                file_format=file_format
             )
 
             return f"Done in {self.format_time(elapsed_time)}! Subtitle file is in the outputs folder.\n\n{subtitle}"
@@ -361,11 +361,11 @@ class WhisperInference(BaseInterface):
             )
 
     @staticmethod
-    def generate_and_write_subtitle(file_name: str,
-                                    transcribed_segments: list,
-                                    add_timestamp: bool,
-                                    subformat: str,
-                                    ) -> str:
+    def generate_and_write_file(file_name: str,
+                                transcribed_segments: list,
+                                add_timestamp: bool,
+                                file_format: str,
+                                ) -> str:
         """
         This method writes subtitle file and returns str to gr.Textbox
         """
@@ -375,13 +375,18 @@ class WhisperInference(BaseInterface):
         else:
             output_path = os.path.join("outputs", f"{file_name}")
 
-        if subformat == "SRT":
-            subtitle = get_srt(transcribed_segments)
-            write_file(subtitle, f"{output_path}.srt")
-        elif subformat == "WebVTT":
-            subtitle = get_vtt(transcribed_segments)
-            write_file(subtitle, f"{output_path}.vtt")
-        return subtitle
+        if file_format == "SRT":
+            content = get_srt(transcribed_segments)
+            write_file(content, f"{output_path}.srt")
+
+        elif file_format == "WebVTT":
+            content = get_vtt(transcribed_segments)
+            write_file(content, f"{output_path}.vtt")
+
+        elif file_format == "txt":
+            content = get_txt(transcribed_segments)
+            write_file(content, f"{output_path}.vtt")
+        return content
 
     @staticmethod
     def format_time(elapsed_time: float) -> str: