diff --git a/gemini-st.py b/gemini-st.py index b33aede..8fb091c 100755 --- a/gemini-st.py +++ b/gemini-st.py @@ -1,5 +1,6 @@ #!/usr/local/lib/pytorch-venv/bin/python3 import os +import argparse from dotenv import load_dotenv import gemini_srt_translator as gst @@ -8,7 +9,7 @@ def create_directory(directory_path): os.makedirs(directory_path, exist_ok=True) -def translate(path, srt_file): +def translate(input_dir, output_dir, srt_file, language): load_dotenv() api_key = os.getenv('GEMINI_API_KEY') if not api_key: @@ -16,21 +17,45 @@ def translate(path, srt_file): return gst.gemini_api_key = api_key - gst.target_language = "Persian" - gst.input_file = os.path.join(path, srt_file) - gst.output_file = os.path.join(path, 'translated', srt_file) + gst.target_language = language + gst.input_file = os.path.join(input_dir, srt_file) + gst.output_file = os.path.join(output_dir, srt_file) gst.model_name = "gemini-2.0-flash" - print(f"Translating {srt_file} ...") + print(f"Translating {srt_file} to {language}...") gst.translate() print("Translation completed successfully.") if __name__ == '__main__': home_dir = os.path.expanduser('~') - directory = os.path.join(home_dir, 'Documents', 'Subtitles') - create_directory(os.path.join(directory, 'translated')) - files = os.listdir(directory) + default_input_dir = os.path.join(home_dir, 'Documents', 'Subtitles') + default_output_dir = os.path.join(default_input_dir, 'translated') + + parser = argparse.ArgumentParser( + description='Translate subtitles using GEMINI SRT translator' + ) + parser.add_argument("-l", + "--language", + help="Target language for translation", + default="Persian", + choices=["English", "Persian", "German", "Italian", "Chinese", "Russian", "Turkish"]) + parser.add_argument("-i", + "--input", + help="Path to input SRT file", + default=default_input_dir) + parser.add_argument("-o", + "--output", + help="Path to output translated SRT file", + default=default_output_dir) + + args = parser.parse_args() + + if args.input == default_input_dir and args.output == default_output_dir: + create_directory(default_output_dir) + files = os.listdir(args.input) srt_files = [f for f in files if f.endswith('.srt')] for file in srt_files: - translate(directory, file) + translate(args.input, args.output, file, args.language) + + print(f"Translated subtitles saved to {args.output}.") \ No newline at end of file