videotrans/task/speech2text.py

import json
import os
import re
import shutil
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import List

from videotrans.configure import contants
from videotrans.configure.config import ROOT_DIR, tr, settings, logger, HOME_DIR
from videotrans.configure import config
from videotrans.recognition import run
from videotrans.task._base import BaseTask
from videotrans.task.taskcfg import TaskCfgSTT



"""
仅语音识别
"""


@dataclass
class SpeechToText(BaseTask):
    cfg: TaskCfgSTT = field(default_factory=TaskCfgSTT, repr=False)
    # 识别后输出的字幕格式，srt txt 等
    out_format: str = field(init=True, default='srt')
    # 在这个子类中，should_recogn 总是 True。
    should_recogn: bool = True
    # 是否需要将生成的字幕复制到原始视频所在目录下，并重命名为视频同名，以方便视频自动加载软字幕
    copysrt_rawvideo: bool = field(default=False, init=True)
    # 存放原始语言字幕
    source_srt_list: List = field(default_factory=list)
    # 插入说话人到字幕开头
    spk_insert: bool = False

    def __post_init__(self):
        super().__post_init__()
        # -1=不启用说话人，0=启用并且不限制说话人数量，>0+1是最大说话人数量
        self.max_speakers = self.cfg.nums_diariz if self.cfg.enable_diariz else -1
        if self.max_speakers > 0:
            self.max_speakers += 1
        # 存放目标文件夹
        if not self.cfg.target_dir:
            self.cfg.target_dir = HOME_DIR + f"/recogn"
        # 转录后的目标字幕文件，先统一转为srt，然后再使用ffmpeg转为其他格式字幕
        self.cfg.target_sub = self.cfg.target_dir + '/' + self.cfg.noextname + '.srt'
        # 临时文件夹
        self.cfg.cache_folder = config.TEMP_DIR + f'/{self.uuid}'
        # 处理为 16k 的wav单通道音频，供模型识别用
        self.cfg.shibie_audio = self.cfg.cache_folder + f'/{self.cfg.noextname}-{time.time()}.wav'
        self.signal(text=tr("Speech Recognition to Word Processing"))

    # 预先处理
    def prepare(self):
        if self._exit(): return
        Path(self.cfg.target_dir).mkdir(parents=True, exist_ok=True)
        Path(self.cfg.cache_folder).mkdir(parents=True, exist_ok=True)
        from videotrans.util.help_ffmpeg import conver_to_16k
        conver_to_16k(self.cfg.name, self.cfg.shibie_audio)

    def recogn(self):
        while 1:
            if self._exit(): return
            # 尚未生成
            if Path(self.cfg.shibie_audio).exists():
                break
            time.sleep(0.5)

        from videotrans.util.help_down import down_file_from_ms
        from videotrans.configure.excepts import SpeechToTextError

        # 需要降噪
        if self.cfg.remove_noise:
            logger.debug('开始降噪')
            from videotrans.process.prepare_audio import remove_noise
            title = tr('Starting to process speech noise reduction, which may take a long time, please be patient')
            down_file_from_ms(f'{ROOT_DIR}/models/onnx', urls=[
                    'https://modelscope.cn/models/himyworld/videotrans/resolve/master/onnx/dpdfnet4.onnx'],
                                        callback=self._process_callback)
            _noise_wav = f"{config.TEMP_DIR}/{self.cfg.noextname}-{os.path.getsize(self.cfg.name)}-removed_noise.wav"
            kw = {
                    "input_file": self.cfg.shibie_audio,
                    "output_file": _noise_wav,
                    "is_cuda": self.cfg.is_cuda
                }
            try:
                _rs = self._new_process(callback=remove_noise, title=title, is_cuda=self.cfg.is_cuda, kwargs=kw)
                if _rs:
                    self.cfg.shibie_audio = _noise_wav
                self.signal(text='remove noise end')
            except Exception as e:
                logger.exception(f'降噪失败，跳过 {e}', exc_info=True)

        if self._exit(): return
        raw_subtitles = run(
            recogn_type=self.cfg.recogn_type,
            uuid=self.uuid,
            model_name=self.cfg.model_name,
            audio_file=self.cfg.shibie_audio,
            detect_language=self.cfg.detect_language,
            cache_folder=self.cfg.cache_folder,
            is_cuda=self.cfg.is_cuda,
            subtitle_type=0,
            max_speakers=self.max_speakers,
            llm_post=self.cfg.rephrase == 1
        )
        if not raw_subtitles or len(raw_subtitles) < 1:
            raise SpeechToTextError(self.cfg.basename + tr('recogn result is empty'))
        self.source_srt_list = raw_subtitles
        self._save_srt_target(self.source_srt_list, self.cfg.target_sub)
        if self._exit() or self.cfg.detect_language == 'auto': return

        # 中英恢复标点符号
        if self.cfg.fix_punc==1 and self.cfg.detect_language[:2] in ['zh', 'en']:
            from videotrans.process.prepare_audio import fix_punc
            down_file_from_ms(f'{ROOT_DIR}/models/puntc', [
                    "https://www.modelscope.cn/models/himyworld/videotrans/resolve/master/puntc/model.onnx",
                    "https://www.modelscope.cn/models/himyworld/videotrans/resolve/master/puntc/config.yaml",
                    "https://www.modelscope.cn/models/himyworld/videotrans/resolve/master/puntc/tokens.json",
            ], callback=self._process_callback)
            text_dict = {f'{it["line"]}': re.sub(r'[,.?!，。？！]', ' ', it["text"]) for it in self.source_srt_list}
            # 序列化后传递文件路径
            text_dict_file=f'{self.cfg.cache_folder}/text_dict_file_{time.time()}.json'
            Path(text_dict_file).write_text(json.dumps(text_dict),encoding="utf-8")
            kw = {"text_dict_file": text_dict_file, "is_cuda": self.cfg.is_cuda}



            try:
                _rs = self._new_process(callback=fix_punc, title=tr("Restoring punct"), kwargs=kw)
                if _rs:
                    text_dict_obj=json.loads(Path(text_dict_file).read_text(encoding='utf-8'))
                    for it in self.source_srt_list:
                        it['text'] = text_dict_obj.get(f'{it["line"]}', it['text'])
                        if self.cfg.detect_language[:2] == 'en':
                            it['text'] = it['text'].replace('，', ',').replace('。', '. ').replace('？', '?').replace(
                                '！', '!')
                    self._save_srt_target(self.source_srt_list, self.cfg.target_sub)
                else:
                    logger.error('标点恢复出错')
            except Exception as e:
                logger.exception(f'恢复标点出错，跳过{e}', exc_info=True)

        # 本身已有说话人识别的，就不再重新断句
        self.signal(text=Path(self.cfg.target_sub).read_text(encoding='utf-8'), type='replace_subtitle')
        if Path(self.cfg.cache_folder + "/speaker.json").exists(): return

        if self.cfg.rephrase == 1:
            # LLM重新断句
            try:
                from videotrans.translator._openaicompat import OpenAICampat
                ob = OpenAICampat(
                    ainame='chatgpt' if settings.get('llm_ai_type', 'chatgpt') != 'deepseek' else 'deepseek',
                    uuid=self.uuid)
                self.signal(text=tr("Re-segmenting..."))
                srt_list = ob.llm_segment(self.source_srt_list)
                if srt_list and len(srt_list) > len(self.source_srt_list) / 2:
                    self.source_srt_list = srt_list
                    self._save_srt_target(self.source_srt_list, self.cfg.target_sub)
                else:
                    logger.error(f'重新断句失败，已恢复原样,原始字幕行:{len(self.source_srt_list)}, 重新断句后字幕行:{len(srt_list)}\n断句结果:\n{srt_list=}')
            except Exception as e:
                self.signal(text=tr("Re-segmenting Error"))
                logger.exception(f"重新断句失败已恢复原样 {e}", exc_info=True)

    def diariz(self):
        if self._exit() or not self.cfg.enable_diariz or Path(self.cfg.cache_folder + "/speaker.json").exists():
            return
        from videotrans.util.help_down import down_file_from_ms, check_and_down_ms
        speaker_type = settings.get('speaker_type', 'built')
        hf_token = settings.get('hf_token')
        if speaker_type == 'built' and self.cfg.detect_language[:2] not in ['zh', 'en']:
            logger.error(f'当前选择 built 说话人分离模型，但不支持当前语言:{self.cfg.detect_language}')
            return
        if speaker_type in ['pyannote', 'reverb'] and not hf_token:
            logger.error(f'当前选择 pyannote 说话人分离模型，但未设置 huggingface.co 的token: {self.cfg.detect_language}')
            return
        hf_endpoit = "https://huggingface.co"
        if speaker_type in ['pyannote', 'reverb']:
            try:
                import requests
                requests.head('https://huggingface.co', timeout=5)
            except Exception:
                logger.exception(f'当前选择 {speaker_type} 说话人分离模型，但无法连接到 https://huggingface.co,可能会失败', exc_info=True)
                hf_endpoit = "https://hf-mirror.com"

        self.precent += 3
        title = tr(f'Begin separating the speakers') + f':{speaker_type}'
        subtitles_file=f'{self.cfg.cache_folder}/diariz-{time.time()}.json'
        Path(subtitles_file).write_text(json.dumps([[it['start_time'], it['end_time']] for it in self.source_srt_list]),encoding='utf-8')
        kw = {
            "input_file": self.cfg.shibie_audio,
            "subtitles_file": subtitles_file,
            "speak_file":self.cfg.cache_folder + "/speaker.json",
            "num_speakers": self.max_speakers,
            "is_cuda": self.cfg.is_cuda
        }
        if speaker_type == 'built':
            down_file_from_ms(f'{ROOT_DIR}/models/onnx', [
                "https://www.modelscope.cn/models/himyworld/videotrans/resolve/master/onnx/seg_model.onnx",
                "https://www.modelscope.cn/models/himyworld/videotrans/resolve/master/onnx/nemo_en_titanet_small.onnx",
                "https://www.modelscope.cn/models/himyworld/videotrans/resolve/master/onnx/3dspeaker_speech_eres2net_large_sv_zh-cn_3dspeaker_16k.onnx"
            ], callback=self._process_callback)
            from videotrans.process.prepare_audio import built_speakers as _run_speakers
            del kw['is_cuda']
            kw['num_speakers'] = -1 if self.max_speakers < 1 else self.max_speakers
            kw['language'] = self.cfg.detect_language
        elif speaker_type == 'ali_CAM':
            check_and_down_ms(model_id='iic/speech_campplus_speaker-diarization_common',
                                    callback=self._process_callback)
            from videotrans.process.prepare_audio import cam_speakers as _run_speakers
        elif speaker_type == 'pyannote':
            from videotrans.process.prepare_audio import pyannote_speakers as _run_speakers
        elif speaker_type == 'reverb':
            from videotrans.process.prepare_audio import reverb_speakers as _run_speakers
        else:
            logger.error(f'当前所选说话人分离模型不支持:{speaker_type=}')
            return
        try:
            if speaker_type in ['pyannote', 'reverb']:
                self.signal(text='Downloading speakers models')
                from huggingface_hub import snapshot_download
                snapshot_download(
                    repo_id="pyannote/speaker-diarization-3.1" if speaker_type == 'pyannote' else "Revai/reverb-diarization-v1",
                    token=hf_token,
                    endpoint=hf_endpoit
                )
            _rs = self._new_process(callback=_run_speakers, title=title,
                                         is_cuda=self.cfg.is_cuda and speaker_type != 'built', kwargs=kw)

            logger.debug('分离说话人成功完成' if _rs else '分离失败说话人失败')
            self.signal(text=tr('separating speakers end'))
        except Exception as e:
            logger.exception(f'说话人分离失败，跳过 {e}', exc_info=True)
        self.signal(text=tr('separating speakers end'))

    def task_done(self):
        if self._exit(): return
        from videotrans.util.help_srt import simple_wrap
        if self.cfg.detect_language and self.cfg.detect_language != 'auto':
            # 处理换行
            maxlen = int(
                settings.get('cjk_len', 15) if self.cfg.detect_language[:2] in contants.CJK_LANG else
                settings.get('other_len', 60))
            for i, it in enumerate(self.source_srt_list):
                it['text'] = simple_wrap(it['text'], maxlen, self.cfg.detect_language)

        # 移除标点符号
        if self.cfg.fix_punc==2:
            from videotrans.util.help_srt import delete_punc
            for i, it in enumerate(self.source_srt_list):
                it['text'] = delete_punc(it['text'])


        if self.cfg.enable_diariz and self.spk_insert and Path(
                self.cfg.cache_folder + "/speaker.json").exists():
            speakers = json.loads(Path(self.cfg.cache_folder + "/speaker.json").read_text(encoding='utf-8'))
            if speakers:
                speakers_len = len(speakers)
                for i, it in enumerate(self.source_srt_list):
                    if i < speakers_len and speakers[i]:
                        it['text'] = f'[{speakers[i]}]{it["text"]}'

        self._save_srt_target(self.source_srt_list, self.cfg.target_sub)
        if self.out_format == 'txt':
            self.cfg.target_sub = self.cfg.target_sub[:-3] + 'txt'
            Path(self.cfg.target_sub).write_text("\r\n".join([it["text"] for it in self.source_srt_list]),
                                                 encoding='utf-8')
        elif self.out_format != 'srt':
            from videotrans.util.help_ffmpeg import runffmpeg
            runffmpeg(['-y', '-i', self.cfg.target_sub, self.cfg.target_sub[:-3] + self.out_format])
            Path(self.cfg.target_sub).unlink(missing_ok=True)
            self.cfg.target_sub = self.cfg.target_sub[:-3] + self.out_format

        if self.copysrt_rawvideo:
            p = Path(self.cfg.name)
            try:
                shutil.copy2(self.cfg.target_sub, f'{p.parent.as_posix()}/{p.stem}.{self.out_format}')
            except shutil.SameFileError:
                pass
        self.set_end(True)