add argparse and minor changes to directories

This commit is contained in:
Hirad 2025-02-28 11:25:34 +03:30
parent 5aafea8e7a
commit 86abc9f58d

View file

@ -1,5 +1,6 @@
#!/usr/local/lib/pytorch-venv/bin/python3 #!/usr/local/lib/pytorch-venv/bin/python3
import os import os
import argparse
from dotenv import load_dotenv from dotenv import load_dotenv
import gemini_srt_translator as gst import gemini_srt_translator as gst
@ -8,7 +9,7 @@ def create_directory(directory_path):
os.makedirs(directory_path, exist_ok=True) os.makedirs(directory_path, exist_ok=True)
def translate(path, srt_file): def translate(input_dir, output_dir, srt_file, language):
load_dotenv() load_dotenv()
api_key = os.getenv('GEMINI_API_KEY') api_key = os.getenv('GEMINI_API_KEY')
if not api_key: if not api_key:
@ -16,21 +17,45 @@ def translate(path, srt_file):
return return
gst.gemini_api_key = api_key gst.gemini_api_key = api_key
gst.target_language = "Persian" gst.target_language = language
gst.input_file = os.path.join(path, srt_file) gst.input_file = os.path.join(input_dir, srt_file)
gst.output_file = os.path.join(path, 'translated', srt_file) gst.output_file = os.path.join(output_dir, srt_file)
gst.model_name = "gemini-2.0-flash" gst.model_name = "gemini-2.0-flash"
print(f"Translating {srt_file} ...") print(f"Translating {srt_file} to {language}...")
gst.translate() gst.translate()
print("Translation completed successfully.") print("Translation completed successfully.")
if __name__ == '__main__': if __name__ == '__main__':
home_dir = os.path.expanduser('~') home_dir = os.path.expanduser('~')
directory = os.path.join(home_dir, 'Documents', 'Subtitles') default_input_dir = os.path.join(home_dir, 'Documents', 'Subtitles')
create_directory(os.path.join(directory, 'translated')) default_output_dir = os.path.join(default_input_dir, 'translated')
files = os.listdir(directory)
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')] srt_files = [f for f in files if f.endswith('.srt')]
for file in srt_files: for file in srt_files:
translate(directory, file) translate(args.input, args.output, file, args.language)
print(f"Translated subtitles saved to {args.output}.")