add argparse and minor changes to directories
This commit is contained in:
parent
5aafea8e7a
commit
86abc9f58d
43
gemini-st.py
43
gemini-st.py
|
@ -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}.")
|
Loading…
Reference in a new issue