611 lines
21 KiB
Python
611 lines
21 KiB
Python
import PIL
|
||
import gradio as gr
|
||
import subprocess
|
||
import os
|
||
import re
|
||
import io
|
||
import base64
|
||
import datetime
|
||
from dotenv import load_dotenv
|
||
from openai import OpenAI
|
||
from PIL import Image
|
||
from docx import Document
|
||
from docx.shared import Pt, Inches
|
||
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
|
||
import markdown
|
||
from html.parser import HTMLParser
|
||
|
||
load_dotenv()
|
||
|
||
|
||
def convert_to_wav(audio_file):
|
||
# 使用 ffmpeg 将音频文件转换为 wav 格式
|
||
output_wav_file = "output.wav"
|
||
print(f"开始转换音频文件 {audio_file} 为 WAV 格式...")
|
||
subprocess.run([
|
||
"ffmpeg", "-y", "-i", audio_file,
|
||
"-ac", "1",
|
||
"-ar", "16000",
|
||
output_wav_file
|
||
])
|
||
print(f"音频文件 {audio_file} 已转换为 WAV 格式，输出文件为 {output_wav_file}")
|
||
return output_wav_file
|
||
|
||
|
||
def transcribe_audio(wav_file, original_filename, offset_time, duration_time):
|
||
output_dir = "/home/tmfc/apps/chyoso-toolkit/whisper_output/"
|
||
# 调用 whisper 命令行程序进行转写
|
||
whisper_cmd = [
|
||
"/home/tmfc/apps/whisper/.venv/bin/whisper",
|
||
"--language", "zh",
|
||
"--output_dir", output_dir,
|
||
"--output_format", "txt",
|
||
"--model", "turbo",
|
||
wav_file
|
||
]
|
||
print(whisper_cmd)
|
||
print(f"开始转写音频文件 {wav_file}...")
|
||
result = subprocess.run(whisper_cmd, capture_output=True, text=True)
|
||
print(result)
|
||
# 将转写结果保存为 txt 文件
|
||
# txt_file = "transcription.txt"
|
||
# with open(txt_file, "w") as f:
|
||
# f.write(result.stdout)
|
||
|
||
txt_file = output_dir + os.path.splitext(original_filename)[0] + ".txt"
|
||
print(f"音频文件 {wav_file} 转写完成，结果已保存为 {txt_file}")
|
||
|
||
return txt_file
|
||
|
||
|
||
def process_audio(audio_file, offset_time, duration_time):
|
||
print("开始处理音频文件...")
|
||
# 获取上传的文件名
|
||
original_filename = os.path.basename(audio_file)
|
||
|
||
# 转换音频文件为 wav 格式
|
||
# wav_file = convert_to_wav(audio_file)
|
||
|
||
# 转写音频文件
|
||
txt_file = transcribe_audio(audio_file, original_filename, offset_time, duration_time)
|
||
|
||
print("音频文件处理完成")
|
||
return txt_file
|
||
|
||
|
||
def direct_transcribe(audio_file, offset_time, duration_time):
|
||
print("开始直接转写音频文件...")
|
||
|
||
# 转写音频文件
|
||
txt_file = transcribe_audio("output.wav", "output.wav", offset_time, duration_time)
|
||
|
||
print("音频文件直接转写完成")
|
||
return txt_file
|
||
|
||
|
||
batch_directory = '/mnt/d/share/audio/'
|
||
|
||
|
||
def list_files():
|
||
# 获取目录下的所有文件
|
||
files = os.listdir(batch_directory)
|
||
# 过滤掉目录，只保留文件
|
||
files = [f for f in files if
|
||
os.path.isfile(os.path.join(batch_directory, f)) and f.lower().endswith(('.mp3', '.m4a'))]
|
||
return files
|
||
|
||
|
||
log_content = ""
|
||
|
||
|
||
def batch_transcribe():
|
||
global log_content
|
||
files = list_files()
|
||
result_file = []
|
||
for file in files:
|
||
# 转换音频文件为 wav 格式
|
||
log_entry = "转换" + file + "为 wav\n"
|
||
log_content += log_entry
|
||
wav_file = convert_to_wav(batch_directory + file)
|
||
log_entry = "转换wav成功，开始转写\n"
|
||
log_content += log_entry
|
||
# 转写音频文件
|
||
txt_file = transcribe_audio(wav_file, file, 10, 0)
|
||
log_entry = "转写 " + file + "完成\n"
|
||
log_content += log_entry
|
||
result_file.append(txt_file)
|
||
return result_file
|
||
|
||
|
||
def display_files():
|
||
files = list_files()
|
||
return "\n".join(files)
|
||
|
||
|
||
def get_log():
|
||
global log_content
|
||
return log_content
|
||
|
||
|
||
def update_log_output():
|
||
return gr.update(value=get_log())
|
||
|
||
|
||
def convert_to_docx(text):
|
||
if text.strip() == "":
|
||
return "输入框不能为空！"
|
||
|
||
# 将输入内容写入 file.md 文件
|
||
with open("file.md", "w") as file:
|
||
file.write(text)
|
||
|
||
# 使用 pandoc 将 file.md 转换为 file.docx
|
||
try:
|
||
subprocess.run(["pandoc", "file.md", "-o", "file.docx"], check=True)
|
||
except subprocess.CalledProcessError as e:
|
||
return f"转换失败: {e}"
|
||
|
||
# 返回 file.docx 文件供用户下载
|
||
return "file.docx"
|
||
|
||
|
||
class MarkdownToDocxParser(HTMLParser):
|
||
"""解析 HTML 并转换为 Word 文档"""
|
||
def __init__(self, document):
|
||
super().__init__()
|
||
self.doc = document
|
||
self.current_paragraph = None
|
||
self.current_run = None
|
||
self.in_bold = False
|
||
self.in_italic = False
|
||
self.in_code = False
|
||
self.in_heading = False
|
||
self.heading_level = 0
|
||
self.list_items = []
|
||
# 表格相关
|
||
self.in_table = False
|
||
self.table_depth = 0 # 追踪表格嵌套深度
|
||
self.current_table = None
|
||
self.current_row = None
|
||
self.current_cell = None
|
||
self.table_rows = []
|
||
self.current_row_cells = []
|
||
self.current_cell_content = [] # 存储单元格内容（包括格式）
|
||
self.is_header_row = False
|
||
|
||
def handle_starttag(self, tag, attrs):
|
||
if tag in ['h1', 'h2', 'h3', 'h4', 'h5', 'h6']:
|
||
self.in_heading = True
|
||
self.heading_level = int(tag[1])
|
||
self.current_paragraph = self.doc.add_heading(level=self.heading_level)
|
||
self.current_paragraph.text = ''
|
||
elif tag == 'p':
|
||
if not self.in_table: # 不在表格中才创建段落
|
||
self.current_paragraph = self.doc.add_paragraph()
|
||
elif self.current_cell is not None:
|
||
# 在表格单元格中，记录换行
|
||
self.current_cell_content.append({'type': 'break'})
|
||
elif tag == 'strong' or tag == 'b':
|
||
self.in_bold = True
|
||
elif tag == 'em' or tag == 'i':
|
||
self.in_italic = True
|
||
elif tag == 'code':
|
||
self.in_code = True
|
||
elif tag == 'li':
|
||
if not self.in_table:
|
||
self.current_paragraph = self.doc.add_paragraph(style='List Bullet')
|
||
elif tag == 'br':
|
||
if self.in_table and self.current_cell is not None:
|
||
self.current_cell_content.append({'type': 'break'})
|
||
elif self.current_paragraph:
|
||
self.current_paragraph.add_run().add_break()
|
||
# 表格处理
|
||
elif tag == 'table':
|
||
self.table_depth += 1
|
||
if self.table_depth == 1: # 只处理最外层表格
|
||
self.in_table = True
|
||
self.table_rows = []
|
||
elif tag == 'thead':
|
||
if self.table_depth == 1:
|
||
self.is_header_row = True
|
||
elif tag == 'tbody':
|
||
if self.table_depth == 1:
|
||
self.is_header_row = False
|
||
elif tag == 'tr':
|
||
if self.table_depth == 1:
|
||
self.current_row_cells = []
|
||
elif tag == 'th' or tag == 'td':
|
||
if self.table_depth == 1:
|
||
self.current_cell = []
|
||
self.current_cell_content = []
|
||
if tag == 'th':
|
||
self.in_bold = True # 表头加粗
|
||
|
||
def handle_endtag(self, tag):
|
||
if tag in ['h1', 'h2', 'h3', 'h4', 'h5', 'h6']:
|
||
self.in_heading = False
|
||
self.heading_level = 0
|
||
elif tag == 'p':
|
||
if not self.in_table:
|
||
self.current_paragraph = None
|
||
elif tag == 'strong' or tag == 'b':
|
||
self.in_bold = False
|
||
elif tag == 'em' or tag == 'i':
|
||
self.in_italic = False
|
||
elif tag == 'code':
|
||
self.in_code = False
|
||
elif tag == 'li':
|
||
if not self.in_table:
|
||
self.current_paragraph = None
|
||
# 表格处理
|
||
elif tag == 'table':
|
||
self.table_depth -= 1
|
||
if self.table_depth == 0:
|
||
self.in_table = False
|
||
self._create_table()
|
||
elif tag == 'thead':
|
||
if self.table_depth == 1:
|
||
self.is_header_row = False
|
||
elif tag == 'tr':
|
||
if self.table_depth == 1 and self.current_row_cells:
|
||
self.table_rows.append(self.current_row_cells)
|
||
self.current_row_cells = []
|
||
elif tag == 'th' or tag == 'td':
|
||
if self.table_depth == 1 and self.current_cell is not None:
|
||
# 保存单元格内容（包含格式信息）
|
||
self.current_row_cells.append(self.current_cell_content.copy())
|
||
self.current_cell = None
|
||
self.current_cell_content = []
|
||
if tag == 'th':
|
||
self.in_bold = False
|
||
|
||
def _create_table(self):
|
||
"""创建 Word 表格"""
|
||
if not self.table_rows:
|
||
return
|
||
|
||
# 计算列数
|
||
max_cols = max(len(row) for row in self.table_rows) if self.table_rows else 0
|
||
if max_cols == 0:
|
||
return
|
||
|
||
# 创建表格
|
||
table = self.doc.add_table(rows=len(self.table_rows), cols=max_cols)
|
||
table.style = 'Light Grid Accent 1'
|
||
|
||
# 填充数据
|
||
for i, row_data in enumerate(self.table_rows):
|
||
row = table.rows[i]
|
||
for j, cell_content_list in enumerate(row_data):
|
||
if j >= len(row.cells):
|
||
continue
|
||
cell = row.cells[j]
|
||
# 清空默认段落
|
||
cell.text = ''
|
||
para = cell.paragraphs[0] if cell.paragraphs else cell.add_paragraph()
|
||
|
||
# 处理单元格内容（支持格式和换行）
|
||
for content_item in cell_content_list:
|
||
if isinstance(content_item, dict):
|
||
if content_item.get('type') == 'break':
|
||
para.add_run().add_break()
|
||
else:
|
||
# 文本内容
|
||
text, is_bold, is_italic = content_item
|
||
run = para.add_run(text)
|
||
if is_bold:
|
||
run.bold = True
|
||
if is_italic:
|
||
run.italic = True
|
||
|
||
# 第一行加粗（表头）
|
||
if i == 0:
|
||
for paragraph in cell.paragraphs:
|
||
for run in paragraph.runs:
|
||
run.bold = True
|
||
|
||
self.table_rows = []
|
||
|
||
def handle_data(self, data):
|
||
if not data.strip() and not self.in_table:
|
||
return
|
||
|
||
# 如果在表格单元格中（只处理最外层表格）
|
||
if self.current_cell is not None and self.table_depth == 1:
|
||
# 保存文本及其格式
|
||
self.current_cell_content.append((data, self.in_bold, self.in_italic))
|
||
return
|
||
|
||
if not self.current_paragraph and not self.in_table:
|
||
self.current_paragraph = self.doc.add_paragraph()
|
||
|
||
if self.current_paragraph:
|
||
run = self.current_paragraph.add_run(data)
|
||
|
||
if self.in_bold:
|
||
run.bold = True
|
||
if self.in_italic:
|
||
run.italic = True
|
||
if self.in_code:
|
||
run.font.name = 'Courier New'
|
||
run.font.size = Pt(10)
|
||
|
||
|
||
def convert_markdown_to_word(markdown_text):
|
||
"""将 Markdown 文本转换为 Word 文档并保存到 word_output 目录（使用纯 Python 实现）"""
|
||
if not markdown_text or markdown_text.strip() == "":
|
||
return "请输入 Markdown 内容！", None
|
||
|
||
# 确保输出目录存在
|
||
output_dir = "word_output"
|
||
if not os.path.exists(output_dir):
|
||
os.makedirs(output_dir)
|
||
|
||
# 生成文件名（使用时间戳）
|
||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
docx_filename = f"document_{timestamp}.docx"
|
||
docx_path = os.path.join(output_dir, docx_filename)
|
||
|
||
try:
|
||
# 将 Markdown 转换为 HTML（启用表格扩展）
|
||
html_content = markdown.markdown(
|
||
markdown_text,
|
||
extensions=['extra', 'nl2br', 'tables']
|
||
)
|
||
|
||
# 创建 Word 文档
|
||
doc = Document()
|
||
|
||
# 解析 HTML 并添加到 Word 文档
|
||
parser = MarkdownToDocxParser(doc)
|
||
parser.feed(html_content)
|
||
|
||
# 保存文档
|
||
doc.save(docx_path)
|
||
message = f"转换成功！文件已保存为: {docx_filename}"
|
||
# 只返回本次生成的文件
|
||
return message, docx_path
|
||
except Exception as e:
|
||
message = f"转换失败: {e}"
|
||
return message, None
|
||
|
||
|
||
def get_word_files():
|
||
"""获取 word_output 目录下的所有 Word 文档"""
|
||
output_dir = "word_output"
|
||
if not os.path.exists(output_dir):
|
||
return []
|
||
|
||
files = []
|
||
for filename in os.listdir(output_dir):
|
||
if filename.endswith('.docx'):
|
||
file_path = os.path.join(output_dir, filename)
|
||
files.append(file_path)
|
||
|
||
# 按修改时间倒序排列
|
||
files.sort(key=lambda x: os.path.getmtime(x), reverse=True)
|
||
return files
|
||
|
||
|
||
def remove_headers(curl_command):
|
||
# 使用正则表达式去除 if-none-match 和 range 标头
|
||
curl_command = re.sub(r'(-H\s*\'if-none-match:[^\\]*\'\s*)', '', curl_command)
|
||
curl_command = re.sub(r'(-H\s*\'range:[^\\]*\'\s*)', '', curl_command)
|
||
return curl_command
|
||
|
||
|
||
def do_download_pdf_file(curl_command, pdf_filename):
|
||
# 去除不需要的标头
|
||
curl_command = remove_headers(curl_command)
|
||
|
||
# 使用 subprocess 调用 curl 命令
|
||
try:
|
||
result = subprocess.run(curl_command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||
# 保存文件到临时文件
|
||
if pdf_filename != "":
|
||
temp_file_path = f"./{pdf_filename}.pdf"
|
||
else:
|
||
temp_file_path = "./download.pdf"
|
||
with open(temp_file_path, 'wb') as f:
|
||
f.write(result.stdout)
|
||
return "File downloaded successfully", temp_file_path
|
||
except subprocess.CalledProcessError as e:
|
||
return f"Failed to download file. Error: {e.stderr.decode()}", None
|
||
|
||
|
||
def download_pdf(curl_command, pdf_filename):
|
||
message, file_path = do_download_pdf_file(curl_command, pdf_filename)
|
||
return file_path
|
||
|
||
|
||
def run_ocr(image):
|
||
# 保存上传的图片到指定路径
|
||
image_path = os.path.expanduser("/home/tmfc/apps/got-ocr/img.png")
|
||
image.save(image_path)
|
||
|
||
# 调用 OCR 命令
|
||
command = [
|
||
"sudo", "-u", "tmfc",
|
||
"/home/tmfc/miniconda3/envs/got/bin/python3", "/home/tmfc/apps/got-ocr/GOT/demo/run_ocr_2.0_crop.py",
|
||
"--model-name", "/home/tmfc/apps/got-ocr/models/",
|
||
"--image-file", image_path
|
||
]
|
||
|
||
out_file = "/home/tmfc/apps/got-ocr/img.txt"
|
||
try:
|
||
with open(out_file, 'w') as f:
|
||
result = subprocess.run(command, stdout=f, stderr=subprocess.PIPE, text=True)
|
||
|
||
except subprocess.CalledProcessError as e:
|
||
return f"识别失败: {e}"
|
||
with open(out_file, 'r', encoding='utf-8') as f:
|
||
content = f.read()
|
||
return content
|
||
|
||
|
||
def ocr_image(image, user_query=""):
|
||
# 压缩图片并降低分辨率
|
||
max_size = 1600
|
||
aspect_ratio = image.width / image.height
|
||
if aspect_ratio > 1: # 宽度大于高度
|
||
new_size = (max_size, int(max_size / aspect_ratio))
|
||
else: # 高度大于或等于宽度
|
||
new_size = (int(max_size * aspect_ratio), max_size)
|
||
compressed_image = image.resize(new_size, PIL.Image.Resampling.LANCZOS)
|
||
buffered = io.BytesIO()
|
||
|
||
compressed_image.save(buffered, format="JPEG", quality=75) # 设置JPEG质量为85
|
||
image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
||
print(f"Base64 字符串长度: {len(image_base64) / 1024:.2f} k")
|
||
messages = [
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "text",
|
||
"text": "识别图片中的文字并以纯文本(txt)格式输出，如果图片分为左右两栏，则先输出左边栏再输出右边栏的内容。" + user_query
|
||
},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/jpeg;base64,{image_base64}"
|
||
}
|
||
}
|
||
]
|
||
}
|
||
]
|
||
client = OpenAI(
|
||
# This is the default and can be omitted
|
||
api_key=os.getenv("GLM_API_KEY"),
|
||
base_url="https://open.bigmodel.cn/api/paas/v4/"
|
||
)
|
||
|
||
chat_completion = client.chat.completions.create(
|
||
messages=messages,
|
||
model="glm-4v-plus",
|
||
)
|
||
return chat_completion.choices[0].message.content
|
||
|
||
|
||
def batch_ocr(files, user_query):
|
||
results = []
|
||
for file in files:
|
||
try:
|
||
image = Image.open(file)
|
||
ocr_result = ocr_image(image, user_query)
|
||
results.append(f"\n\n{ocr_result}\n\n")
|
||
except Exception as e:
|
||
print(e)
|
||
|
||
# 将结果写入 Markdown 文件
|
||
output_file = "ocr_results.txt"
|
||
with open(output_file, "w") as f:
|
||
f.write("\n".join(results))
|
||
|
||
return output_file
|
||
|
||
|
||
with gr.Blocks() as iface:
|
||
gr.Markdown("# 大模型工具集")
|
||
with gr.Tabs():
|
||
with gr.TabItem("音频转写"):
|
||
with gr.Row():
|
||
audio_input = gr.Audio(type="filepath", label="上传音频文件")
|
||
with gr.Column():
|
||
offset_input = gr.Number(label="偏移时间 (秒)")
|
||
duration_input = gr.Number(label="转写时长 (秒)")
|
||
|
||
with gr.Row():
|
||
process_button = gr.Button("处理并转写")
|
||
direct_transcribe_button = gr.Button("直接转写")
|
||
|
||
output_file = gr.File(label="转写结果")
|
||
|
||
process_button.click(process_audio, inputs=[audio_input, offset_input, duration_input], outputs=output_file)
|
||
direct_transcribe_button.click(direct_transcribe, inputs=[audio_input, offset_input, duration_input],
|
||
outputs=output_file)
|
||
with gr.TabItem("音频批量转写"):
|
||
with gr.Row():
|
||
link_upload = gr.HTML(value='<a href="https://webd.willking.tech" target="_blank">点击上传文件</a>')
|
||
with gr.Row():
|
||
file_list = gr.Textbox(label="文件列表")
|
||
with gr.Column():
|
||
list_file_button = gr.Button("刷新文件")
|
||
batch_process_button = gr.Button("批量处理")
|
||
with gr.Column():
|
||
batch_output_file = gr.Files(label="批量转写结果")
|
||
|
||
with gr.Row():
|
||
log_output = gr.Textbox(label="日志信息", lines=10)
|
||
|
||
list_file_button.click(fn=display_files, outputs=file_list)
|
||
|
||
batch_process_button.click(batch_transcribe, outputs=batch_output_file)
|
||
|
||
with gr.Tab("Markdown 转 Word"):
|
||
gr.Markdown("## Markdown 转 Word 转换器")
|
||
with gr.Row():
|
||
markdown_input = gr.Textbox(
|
||
lines=15,
|
||
placeholder="请在此输入 Markdown 内容...",
|
||
label="Markdown 内容"
|
||
)
|
||
with gr.Row():
|
||
convert_md_button = gr.Button("转换为 Word", variant="primary")
|
||
with gr.Row():
|
||
conversion_status = gr.Textbox(label="转换状态", interactive=False)
|
||
with gr.Row():
|
||
word_file_output = gr.File(label="生成的 Word 文档")
|
||
|
||
convert_md_button.click(
|
||
convert_markdown_to_word,
|
||
inputs=markdown_input,
|
||
outputs=[conversion_status, word_file_output]
|
||
)
|
||
|
||
with gr.Tab("下载pdf"):
|
||
gr.Markdown("## pdf 下载指令修复")
|
||
with gr.Row():
|
||
curl_text_input = gr.Textbox(lines=10, placeholder="请在此输入cURL脚本...")
|
||
with gr.Row():
|
||
with gr.Column():
|
||
pdf_filename = gr.Textbox(placeholder="输入文件名")
|
||
with gr.Column():
|
||
pdf_download_button = gr.Button("下载")
|
||
with gr.Row():
|
||
pdf_download_result = gr.File(label="下载pdf文件")
|
||
|
||
pdf_download_button.click(download_pdf, inputs=[curl_text_input, pdf_filename], outputs=pdf_download_result)
|
||
|
||
with gr.Tab("图片识别"):
|
||
gr.Markdown("## OCR 图片识别")
|
||
with gr.Row():
|
||
with gr.Column():
|
||
image_input = gr.Image(type="pil", label="上传图片")
|
||
with gr.Column():
|
||
btn_recognize = gr.Button("识别")
|
||
text_output = gr.Textbox(label="OCR 识别结果")
|
||
|
||
btn_recognize.click(fn=run_ocr, inputs=image_input, outputs=text_output)
|
||
|
||
with gr.Tab("批量图片识别"):
|
||
gr.Markdown("## OCR 图片批量识别")
|
||
with gr.Row():
|
||
with gr.Column():
|
||
input_files = gr.File(file_count="multiple", label="上传图片")
|
||
with gr.Column():
|
||
output_file = gr.File(label="下载识别结果")
|
||
with gr.Row():
|
||
process_button = gr.Button("开始识别")
|
||
with gr.Row():
|
||
user_query_text = gr.Textbox(label="额外识别要求", placeholder="输入额外的要求")
|
||
|
||
# 绑定按钮点击事件
|
||
process_button.click(batch_ocr, inputs=[input_files, user_query_text], outputs=output_file)
|
||
|
||
# 使用 Gradio 5 的定时器 API 周期性更新日志输出
|
||
timer = gr.Timer(1.0)
|
||
timer.tick(fn=update_log_output, outputs=[log_output])
|
||
|
||
iface.launch(server_name="0.0.0.0", server_port=7861)
|