[tool.poetry]name = "upbox"version = "0.1.0"description = ""authors = ["rocksun <daijun@gmail.com>"]readme = "README.md"
[tool.poetry.dependencies]python = "^3.10"ffmpeg-python = "^0.2.0"llama-index = "^0.10.25"llama-index-llms-gemini = "^0.1.6"pysubs2 = "^1.6.1"# yt-dlp = "^2024.4.9"# typer = "^0.12.3"# faster-whisper = "^1.0.1"

[build-system]requires = ["poetry-core"]build-backend = "poetry.core.masonry.api"

ffmpeg -i my_file.mkv outfile.vtt

def _guess_eng_subtitle_index(video_path):probe = ffmpeg.probe(video_path)streams = probe['streams']for index, stream in enumerate(streams):if stream.get('codec_type') == 'subtitle' and stream.get('tags', {}).get('language') == 'eng':return indexfor index, stream in enumerate(streams):if stream['codec_type'] == 'subtitle' and stream.get('tags', {}).get('title', "").lower().find("english")!=-1 :return indexreturn -1
def _extract_subtitle_by_index(video_path, output_path, index):return ffmpeg.input(video_path).output(output_path, map='0:'+str(index)).run()
def extract_subtitle(video_path, en_subtitle_path):# get the streams from video with ffprobeindex = _guess_eng_subtitle_index(video_path)if index == -1:return -1return _extract_subtitle_by_index(video_path, en_subtitle_path, index)

class UpSubs:def __init__(self, subs_path):self.subs = pysubs2.load(subs_path)
def get_subtitle_text(self):text = ""for sub in self.subs:text += sub.text + "\n\n"return text
def get_subtitle_text_with_index(self):text = ""for i, sub in enumerate(self.subs):text += "chunk-"+str(i) + ":\n" + sub.text.replace("\\N", " ") + "\n\n"return textdef save(self, output_path):self.subs.save(output_path)
def clean(self):indexes = []for i, sub in enumerate(self.subs):# remove xml tag and line change in sub textsub.text = re.sub(r"<[^>]+>", "", sub.text)sub.text = sub.text.replace("\\N", " ")
def fill(self, text):text = text.strip()pattern = r"\n\s*\n"paragraphs = re.split(pattern, text)for para in paragraphs:try:firtline = para.split("\n")[0]countstr = firtline[6:len(firtline)-1]# print(countstr)index = int(countstr)p = "\n".join(para.split("\n")[1:])self.subs[index].text = pexcept Exception as e:print(f"Error merge paragraph : \n {para} \n with exception: \n {e}")raise(e)def merge_dual(self, subspath):second_subs = pysubs2.load(subspath)merged_subs = SSAFile()if len(self.subs.events) == len(second_subs.events):for i, first_event in enumerate(self.subs.events):second_event = second_subs[i]if first_event.text == second_event.text:merged_event = SSAEvent(first_event.start, first_event.end, first_event.text)else:merged_event = SSAEvent(first_event.start, first_event.end, first_event.text + '\n' + second_event.text)merged_subs.append(merged_event)return merged_subsreturn None

1200:02:30,776 --> 00:02:34,780Not even the great Dragon Warrior.
1300:02:43,830 --> 00:02:45,749Oh, where is Po?
1400:02:45,749 --> 00:02:48,502He was supposed to be here hours ago.

chunk-12Not even the great Dragon Warrior.
chunk-13Oh, where is Po?
chunk-14He was supposed to be here hours ago.

def complete(prompt, max_tokens=32760):prompt = prompt.strip()if not prompt:return ""safety_settings = [{"category": "HARM_CATEGORY_HARASSMENT","threshold": "BLOCK_NONE"},{"category": "HARM_CATEGORY_HATE_SPEECH","threshold": "BLOCK_NONE"},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","threshold": "BLOCK_NONE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","threshold": "BLOCK_NONE"},]
retries = 3for _ in range(retries):try:return Gemini(max_tokens=max_tokens, safety_settings=safety_settings, temperature = 0.01).complete(prompt).textexcept Exception as e:print(f"Error completing prompt: {prompt} \n with error: \n ")traceback.print_exc()return ""

http_proxy=http://192.168.0.107:7890https_proxy=http://192.168.0.107:7890GOOGLE_API_KEY=[your-api-key]

def tran_subtitles(fixed_subtitle, zh_subtitle=None, cncf = False, chunk_size=3000):subtitle_base = os.path.splitext(fixed_subtitle)[0]video_base = os.path.splitext(subtitle_base)[0]if zh_subtitle is None:zh_subtitle = video_base + ".zh-fixed.vtt"if os.path.exists(zh_subtitle):print(f"zh subtitle {zh_subtitle} already translated, skip to translate.")return 1
prompt_tpl = MOVIE_TRAN_PROMPT_TPLopts = { }
srtp = UpSubs(fixed_subtitle)text = srtp.get_subtitle_text_with_index()
process_text(srtp, text, prompt_tpl, opts, chunk_size = chunk_size)srtp.save(zh_subtitle)
return zh_subtitle

MOVIE_TRAN_PROMPT_TPL = """你是个专业电影字幕翻译，你需要将一份英文字幕翻译成中文。[需要翻译的英文字幕]:
{content}
# [中文字幕]:"""

def process_text(subs, text, prompt_tpl, opts, chunk_size=2500):# ret = ""chunks = _split_subtitles(text, chunk_size)for(i, chunk) in enumerate(chunks):print("process chunk ({}/{})".format(i+1,len(chunks)))# if i==4:# break# format string with all the field in a dict opts["content"] = chunkprompt = prompt_tpl.format(**opts)
print(prompt)out = complete(prompt, max_tokens=32760)subs.fill(out)print(out)

1000:02:22,184 --> 00:02:27,606Let it be known from the highest mountainto the lowest valley that Tai Lung lives,
1100:02:27,606 --> 00:02:30,776and no one will stand in his way.
1200:02:30,776 --> 00:02:34,780Not even the great Dragon Warrior.
1300:02:43,830 --> 00:02:45,749Oh, where is Po?

1000:02:22,184 --> 00:02:27,606让最高的山峰和最低的山谷都知道，泰隆还活着，
1100:02:27,606 --> 00:02:30,776没人能阻挡他。
1200:02:30,776 --> 00:02:34,780即使是伟大的神龙大侠也不行。
1300:02:43,830 --> 00:02:45,749哦，阿宝在哪儿？