275 lines
13 KiB
Python
275 lines
13 KiB
Python
# filename: srt_refiner_v4.2.py
|
||
|
||
import tkinter as tk
|
||
from tkinter import filedialog, ttk, scrolledtext, messagebox
|
||
import sv_ttk
|
||
import threading
|
||
import queue
|
||
import os
|
||
import requests # <--- 核心改动：使用 requests
|
||
import json # <--- 核心改动：用于处理JSON数据
|
||
import srt
|
||
import re
|
||
from datetime import timedelta
|
||
|
||
# --- 核心配置 ---
|
||
OLLAMA_HOST = "http://127.0.0.1:11434"
|
||
CONTEXT_WINDOW = 2
|
||
|
||
# --- GUI 应用 ---
|
||
class App:
|
||
def __init__(self, root):
|
||
self.root = root
|
||
self.root.title("字幕内容精炼师 V4.2 - Requests直连版")
|
||
self.root.geometry("800x600")
|
||
|
||
# --- 核心改动：移除 ollama.Client ---
|
||
self.is_ollama_available = False
|
||
|
||
self.srt_path = ""
|
||
self.original_subs = []
|
||
self.refined_subs = []
|
||
self.processing_thread = None
|
||
self.gui_queue = queue.Queue()
|
||
|
||
self.vars = {
|
||
"srt_path": tk.StringVar(),
|
||
"ollama_model": tk.StringVar(),
|
||
"status": tk.StringVar(value="准备就绪")
|
||
}
|
||
|
||
self.build_ui()
|
||
sv_ttk.set_theme("dark")
|
||
|
||
self.root.after(100, self.load_ollama_models)
|
||
self.root.after(100, self.process_queue)
|
||
|
||
def build_ui(self):
|
||
# ... (UI部分完全不变)
|
||
main_frame = ttk.Frame(self.root, padding=10)
|
||
main_frame.pack(fill=tk.BOTH, expand=True)
|
||
main_frame.rowconfigure(3, weight=1)
|
||
main_frame.columnconfigure(0, weight=1)
|
||
f1 = ttk.Frame(main_frame)
|
||
f1.grid(row=0, column=0, sticky="ew", pady=(0, 10))
|
||
f1.columnconfigure(1, weight=1)
|
||
ttk.Label(f1, text="SRT文件:").grid(row=0, column=0, padx=(0, 5))
|
||
ttk.Entry(f1, textvariable=self.vars['srt_path'], state="readonly").grid(row=0, column=1, sticky="ew", padx=(0, 5))
|
||
ttk.Button(f1, text="选择...", command=self.select_srt_file).grid(row=0, column=2)
|
||
f2 = ttk.Frame(main_frame)
|
||
f2.grid(row=1, column=0, sticky="ew", pady=(0, 10))
|
||
f2.columnconfigure(1, weight=1)
|
||
ttk.Label(f2, text="Ollama模型:").grid(row=0, column=0, padx=(0, 5))
|
||
self.model_combo = ttk.Combobox(f2, textvariable=self.vars['ollama_model'], state="readonly")
|
||
self.model_combo.grid(row=0, column=1, sticky="ew", padx=(0, 5))
|
||
f3 = ttk.Frame(main_frame)
|
||
f3.grid(row=2, column=0, sticky="ew", pady=(0, 15))
|
||
self.start_button = ttk.Button(f3, text="开始精炼", style="Accent.TButton", command=self.start_processing, state="disabled")
|
||
self.start_button.pack(side=tk.LEFT, padx=(0, 5))
|
||
self.save_button = ttk.Button(f3, text="另存为...", command=self.save_srt, state="disabled")
|
||
self.save_button.pack(side=tk.LEFT)
|
||
log_frame = ttk.Labelframe(main_frame, text="处理日志")
|
||
log_frame.grid(row=3, column=0, sticky="nsew")
|
||
log_frame.rowconfigure(0, weight=1)
|
||
log_frame.columnconfigure(0, weight=1)
|
||
self.log_text = scrolledtext.ScrolledText(log_frame, wrap=tk.WORD, state="disabled", height=10)
|
||
self.log_text.grid(row=0, column=0, sticky="nsew", padx=5, pady=5)
|
||
status_bar = ttk.Frame(main_frame)
|
||
status_bar.grid(row=4, column=0, sticky="ew", pady=(5, 0))
|
||
self.progress_bar = ttk.Progressbar(status_bar, orient=tk.HORIZONTAL, mode='determinate')
|
||
self.progress_bar.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 10))
|
||
ttk.Label(status_bar, textvariable=self.vars['status']).pack(side=tk.LEFT)
|
||
|
||
def select_srt_file(self):
|
||
path = filedialog.askopenfilename(filetypes=[("SRT Subtitles", "*.srt")])
|
||
if not path: return
|
||
try:
|
||
with open(path, 'r', encoding='utf-8-sig') as f: content = f.read()
|
||
self.original_subs = list(srt.parse(content))
|
||
self.srt_path = path
|
||
self.vars['srt_path'].set(os.path.basename(path))
|
||
self.log_message(f"已加载文件: {os.path.basename(path)}, 共 {len(self.original_subs)} 条字幕。")
|
||
if self.is_ollama_available:
|
||
self.start_button.config(state="normal")
|
||
self.save_button.config(state="disabled")
|
||
self.refined_subs = []
|
||
except Exception as e:
|
||
messagebox.showerror("加载失败", f"无法加载或解析文件: {e}")
|
||
|
||
def save_srt(self):
|
||
# ... (此方法不变)
|
||
if not self.refined_subs:
|
||
messagebox.showwarning("无内容", "没有可保存的精炼后字幕。")
|
||
return
|
||
original_basename = os.path.splitext(os.path.basename(self.srt_path))[0]
|
||
save_path = filedialog.asksaveasfilename(
|
||
defaultextension=".srt", initialfile=f"{original_basename}_refined.srt",
|
||
filetypes=[("SRT Subtitles", "*.srt")]
|
||
)
|
||
if not save_path: return
|
||
try:
|
||
content_to_save = srt.compose(self.refined_subs)
|
||
with open(save_path, 'w', encoding='utf-8') as f:
|
||
f.write(content_to_save)
|
||
messagebox.showinfo("保存成功", f"精炼后的SRT文件已保存至:\n{save_path}")
|
||
except Exception as e:
|
||
messagebox.showerror("保存失败", f"无法保存文件: {e}")
|
||
|
||
def start_processing(self):
|
||
# ... (此方法不变)
|
||
if not self.is_ollama_available:
|
||
messagebox.showerror("错误", "Ollama 服务未连接，无法开始处理。")
|
||
return
|
||
if not self.original_subs:
|
||
messagebox.showerror("错误", "请先加载一个SRT文件。")
|
||
return
|
||
if not self.vars['ollama_model'].get():
|
||
messagebox.showerror("错误", "请选择一个Ollama模型。")
|
||
return
|
||
self.start_button.config(state="disabled")
|
||
self.save_button.config(state="disabled")
|
||
self.log_text.config(state="normal"); self.log_text.delete(1.0, tk.END); self.log_text.config(state="disabled")
|
||
self.refined_subs = []
|
||
self.progress_bar['value'] = 0
|
||
self.progress_bar['maximum'] = len(self.original_subs)
|
||
self.processing_thread = threading.Thread(
|
||
target=self.refine_worker,
|
||
args=(list(self.original_subs), self.vars['ollama_model'].get()), daemon=True
|
||
)
|
||
self.processing_thread.start()
|
||
|
||
def log_message(self, msg):
|
||
self.gui_queue.put({"type": "log", "data": msg})
|
||
|
||
def process_queue(self):
|
||
# ... (此方法不变)
|
||
try:
|
||
while True:
|
||
msg = self.gui_queue.get_nowait()
|
||
msg_type = msg.get("type")
|
||
if msg_type == "log":
|
||
self.log_text.config(state="normal")
|
||
self.log_text.insert(tk.END, msg["data"] + "\n")
|
||
self.log_text.see(tk.END)
|
||
self.log_text.config(state="disabled")
|
||
elif msg_type == "error":
|
||
self.log_message(f"错误: {msg['data']}")
|
||
messagebox.showerror("Ollama 错误", msg['data'])
|
||
elif msg_type == "models_loaded":
|
||
models = msg['data']
|
||
if models:
|
||
self.model_combo['values'] = models
|
||
self.model_combo.set(models[0])
|
||
self.log_message(f"成功检测到模型: {', '.join(models)}")
|
||
self.is_ollama_available = True
|
||
if self.original_subs: self.start_button.config(state="normal")
|
||
else:
|
||
self.log_message("未检测到任何本地Ollama模型。")
|
||
elif msg_type == "progress":
|
||
self.progress_bar['value'] = msg['current']
|
||
self.vars['status'].set(f"处理中... {msg['current']}/{msg['total']}")
|
||
elif msg_type == "result":
|
||
self.refined_subs.append(msg['data'])
|
||
elif msg_type == "finish":
|
||
self.start_button.config(state="normal")
|
||
self.save_button.config(state="normal")
|
||
self.vars['status'].set("精炼完成！")
|
||
self.log_message("\n--- 所有字幕精炼完成！---")
|
||
messagebox.showinfo("完成", "所有字幕已成功精炼！")
|
||
except queue.Empty: pass
|
||
finally: self.root.after(100, self.process_queue)
|
||
|
||
# --- 核心改动：使用 requests 获取模型列表 ---
|
||
def load_ollama_models(self):
|
||
self.log_message("正在连接Ollama并获取模型列表...")
|
||
def _load():
|
||
try:
|
||
# Ollama列出模型的API端点是 /api/tags
|
||
response = requests.get(f"{OLLAMA_HOST}/api/tags", timeout=5)
|
||
response.raise_for_status() # 如果状态码不是2xx，则抛出异常
|
||
|
||
models_data = response.json()['models']
|
||
model_names = [m['name'] for m in models_data]
|
||
|
||
self.gui_queue.put({"type": "models_loaded", "data": model_names})
|
||
except requests.exceptions.RequestException as e:
|
||
self.gui_queue.put({"type": "error", "data": f"无法连接到Ollama服务: {e}\n请确保Ollama正在运行并且地址({OLLAMA_HOST})正确。"})
|
||
except (KeyError, IndexError) as e:
|
||
self.gui_queue.put({"type": "error", "data": f"解析Ollama模型列表时出错: {e}\n返回的数据格式可能不正确。"})
|
||
threading.Thread(target=_load, daemon=True).start()
|
||
|
||
# --- 核心改动：使用 requests 调用生成API ---
|
||
def refine_worker(self, subs_to_process, model_name):
|
||
total_subs = len(subs_to_process)
|
||
for i, current_sub in enumerate(subs_to_process):
|
||
start_index = max(0, i - CONTEXT_WINDOW)
|
||
end_index = min(total_subs, i + 1 + CONTEXT_WINDOW)
|
||
context_block = subs_to_process[start_index:end_index]
|
||
|
||
prompt = self.build_prompt(context_block, i - start_index)
|
||
self.log_message(f"\n[正在处理 {i+1}/{total_subs}] 原文: {current_sub.content}")
|
||
|
||
# 构造请求体
|
||
payload = {
|
||
"model": model_name,
|
||
"prompt": prompt,
|
||
"stream": False,
|
||
"options": {'temperature': 0.2, 'num_predict': 128}
|
||
}
|
||
|
||
try:
|
||
# Ollama生成的API端点是 /api/generate
|
||
response = requests.post(f"{OLLAMA_HOST}/api/generate", json=payload, timeout=60)
|
||
response.raise_for_status()
|
||
|
||
response_data = response.json()
|
||
refined_text = response_data['response'].strip().replace("\n", " ")
|
||
refined_text = re.sub(r'^["\'“‘]|["\'”’]$', '', refined_text)
|
||
|
||
self.log_message(f"-> 精炼后: {refined_text}")
|
||
|
||
new_sub = srt.Subtitle(
|
||
index=current_sub.index,
|
||
start=current_sub.start,
|
||
end=current_sub.end,
|
||
content=refined_text
|
||
)
|
||
self.gui_queue.put({"type": "result", "data": new_sub})
|
||
except Exception as e:
|
||
self.log_message(f"!! 警告: 处理第 {i+1} 条时出错: {e}。将使用原始内容。")
|
||
self.gui_queue.put({"type": "result", "data": current_sub})
|
||
|
||
self.gui_queue.put({"type": "progress", "current": i + 1, "total": total_subs})
|
||
self.gui_queue.put({"type": "finish"})
|
||
|
||
def build_prompt(self, context_block, current_index_in_block):
|
||
# ... (此方法不变)
|
||
prompt_template = """你是一个专业的视频字幕精炼师。任务是优化“待处理字幕”，使其更适合专业配音。
|
||
规则：
|
||
1. 改为流畅、专业的书面语，但【重要】必须保留所有的核心操作指令和细节。
|
||
2. 【优先】去除明显的口语化词汇（如'嗯'、'啊'）、重复和不必要的填充词（如'然后'、'就是说'）。
|
||
3. 【次要】在不影响信息完整性的前提下，可以适当缩短句子。
|
||
4. 【重要】只输出精炼后的字幕文本，不要包含任何标签、解释或引号。。
|
||
---
|
||
[上下文]
|
||
{context_text}
|
||
---
|
||
[待处理字幕]
|
||
{target_text}
|
||
---
|
||
[精炼后的文本]："""
|
||
context_lines = []
|
||
target_text = ""
|
||
for j, sub in enumerate(context_block):
|
||
clean_content = sub.content.replace('\n', ' ').strip()
|
||
if j == current_index_in_block:
|
||
target_text = clean_content
|
||
context_lines.append(f"-> {clean_content} <- (待处理)")
|
||
else:
|
||
context_lines.append(clean_content)
|
||
return prompt_template.format(context_text="\n".join(context_lines), target_text=target_text)
|
||
|
||
if __name__ == "__main__":
|
||
root = tk.Tk()
|
||
app = App(root)
|
||
root.mainloop() |