wmdigit.py

import copy
import glob
import logging
import os
import time

from . import wmdigit_cut, wmdigit_transcribe, utils


class Wmdigit:
    def __init__(self, args):
        self.args = args

    def run(self):
        assert len(self.args.inputs) == 1, "Must provide a single file"
        self._pipeline()

    def _pipeline(self):
        media_file = self.args.inputs[0]
        assert utils.is_video(media_file), "Must provide a video file"

        args = copy.deepcopy(self.args)

        srt_fn = utils.change_ext(media_file, "srt")
        md_fn = utils.change_ext(media_file, "md")
        # 1、视频生成srt和md
        args.inputs = [media_file]
        # 如果目标语言不是中文，则提示whisper翻译全部字幕
        if args.lang and args.lang != "zh":
            args.prompt = f"Subtitles must be fully translated into {args.lang}"
        logging.info(f"Transcribe {media_file} lang={args.lang} promt={args.prompt}")
        wmdigit_transcribe.Transcribe(args).run()

        # 2、从md生成cut视频
        args.inputs = [media_file, md_fn, srt_fn]
        wmdigit_cut.Cutter(args).run()