add ffmpeg

.gitignore

@@ -1,3 +1,4 @@
 /build/
 /dist/
 /venv/
+__pycache__

config.py

@@ -20,6 +20,12 @@ VIDEO_RESOLUTION = "1280x720"
 # [ffmpeg]
 # exec
 FFMPEG_EXEC = "ffmpeg"
+# nvidia_gpu
+FFMPEG_USE_NVIDIA_GPU = False
+# intel_gpu
+FFMPEG_USE_INTEL_GPU = False
+# vaapi
+FFMPEG_USE_VAAPI = False
 # [handbrake]
 # exec
 HANDBRAKE_EXEC = "HandBrakeCli"
@@ -91,8 +97,11 @@ def load_config():
         VIDEO_CLIP_OVERFLOW_SEC = section.getfloat('overflow_sec', VIDEO_CLIP_OVERFLOW_SEC)
     if config.has_section("ffmpeg"):
         section = config['ffmpeg']
-        global FFMPEG_EXEC
+        global FFMPEG_EXEC, FFMPEG_USE_NVIDIA_GPU, FFMPEG_USE_INTEL_GPU, FFMPEG_USE_VAAPI
         FFMPEG_EXEC = section.get('exec', FFMPEG_EXEC)
+        FFMPEG_USE_NVIDIA_GPU = section.getboolean('nvidia_gpu', FFMPEG_USE_NVIDIA_GPU)
+        FFMPEG_USE_INTEL_GPU = section.getboolean('intel_gpu', FFMPEG_USE_INTEL_GPU)
+        FFMPEG_USE_VAAPI = section.getboolean('vaapi', FFMPEG_USE_VAAPI)
     if config.has_section("handbrake"):
         section = config['handbrake']
         global HANDBRAKE_EXEC, HANDBRAKE_PRESET_FILE, HANDBRAKE_PRESET, HANDBRAKE_ENCOPT

danmaku_workflow_ffmpeg.py

@@ -0,0 +1,557 @@
+# 工作流
+import os.path
+import platform
+import subprocess
+import sys
+import traceback
+from hashlib import md5
+from typing import Optional, IO, Union
+
+from datetime import datetime, timedelta
+from PyQt5 import QtGui
+from PyQt5.QtCore import Qt, QThread, pyqtSignal
+from PyQt5.QtWidgets import QWidget, QLabel, QApplication, QFrame, QVBoxLayout, QPushButton, \
+    QSizePolicy, QMessageBox
+from danmaku_xml_helper import get_file_start, diff_danmaku_files, NoDanmakuException
+from config import load_config, \
+    DANMAKU_EXEC, DANMAKU_SPEED, DANMAKU_FONT_NAME, DANMAKU_FONT_SIZE, \
+    VIDEO_CLIP_EACH_SEC, VIDEO_CLIP_OVERFLOW_SEC, VIDEO_RESOLUTION, VIDEO_OUTPUT_DIR, \
+    FFMPEG_EXEC, FFMPEG_USE_INTEL_GPU, FFMPEG_USE_NVIDIA_GPU
+    
+
+
+class Job:
+    DANMAKU_ENCODE = 0
+    PURE_SPLIT = 1
+
+    def __init__(self):
+        super(Job, self).__init__()
+        self.video = None
+        self.type = self.PURE_SPLIT
+        self.danmaku: list[str] = []
+        self.subtitles: list[str] = []
+
+    def __repr__(self):
+        return "Job对象：Video[{}];Type[{}]".format(self.video, self.type)
+
+
+class WorkLabel(QLabel):
+    def __init__(self, *args, **kwargs):
+        super(WorkLabel, self).__init__(*args, **kwargs)
+        self.workVideo = None
+        self.workDanmaku: list[str] = []
+        self.running: bool = False
+        self.finished: bool = False
+        self.init_ui()
+
+    def _update_text(self):
+        if self.workVideo is None:
+            text = "请拖入视频"
+        else:
+            text = "工作视频：{}".format(self.workVideo)
+        if len(self.workDanmaku) == 0:
+            self.setText("{}\n分割视频模式".format(text))
+        else:
+            self.setText("{}\n压制弹幕模式\n基准弹幕：{}".format(text, "\n附加弹幕：".join(self.workDanmaku)))
+        self.adjustSize()
+
+    def mouseDoubleClickEvent(self, a0: QtGui.QMouseEvent) -> None:
+        if self.finished:
+            self.finished = False
+        elif self.workVideo is None:
+            self.parent().labelDestroy.emit(self)
+            return self.deleteLater()
+        else:
+            self.workVideo = None
+            self.workDanmaku = []
+        self.init_ui()
+
+    def init_ui(self):
+        font = QtGui.QFont()
+        font.setPixelSize(12)
+        self.setFont(font)
+        self.setFrameShape(QFrame.Box)
+        self.setStyleSheet("""QLabel {
+            background-color: white;
+        }""")
+        self.setAlignment(Qt.AlignTop)
+        self.setWordWrap(True)
+        self.setMinimumHeight(64)
+        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
+        self._update_text()
+        self.setAcceptDrops(True)
+
+    def dragEnterEvent(self, a0: QtGui.QDragEnterEvent) -> None:
+        if a0.mimeData().hasUrls():
+            a0.accept()
+        else:
+            a0.ignore()
+        print("Enter Label", a0.mimeData().urls())
+
+    def dropEvent(self, a0: QtGui.QDropEvent) -> None:
+        urls = a0.mimeData().urls()
+        for url in urls:
+            # 判断是否为视频
+            path = url.path().strip("/")
+            if path.endswith(".mp4") or path.endswith(".flv"):
+                self.set_video(path)
+            elif path.endswith(".xml"):
+                self.add_danmaku(path)
+            else:
+                print("Unknown File", path)
+        self._update_text()
+
+    def set_video(self, file):
+        if not os.path.isfile(file):
+            raise FileNotFoundError(file)
+        if file == self.workVideo:
+            return print("Same Video", file)
+        self.workVideo = file
+        self._update_text()
+
+    def add_danmaku(self, file: str) -> None:
+        if not os.path.isfile(file):
+            raise FileNotFoundError(file)
+        if file in self.workDanmaku:
+            return print("Already Added File", file)
+        self.workDanmaku.append(file)
+        self._update_text()
+
+    def have_job(self) -> bool:
+        return self.workVideo is not None and not self.running and not self.finished
+
+    def get_job(self) -> Job:
+        job = Job()
+        job.video = self.workVideo
+        if len(self.workDanmaku) > 0:
+            job.type = Job.DANMAKU_ENCODE
+            job.danmaku = self.workDanmaku
+        else:
+            job.type = Job.PURE_SPLIT
+            job.danmaku = []
+        return job
+
+    def start_running(self) -> None:
+        self.running = True
+        self.finished = False
+        self.setStyleSheet("""QLabel {
+            color: white;
+            background-color: red;
+        }""")
+
+    def stop_running(self) -> None:
+        self.running = False
+        self.finished = True
+        self.setStyleSheet("""QLabel {
+            color: white;
+            background-color: green;
+        }""")
+
+
+class HomePage(QWidget):
+    showMessageBox = pyqtSignal(str, str)
+    labelDestroy = pyqtSignal(WorkLabel)
+    processCurTime = pyqtSignal(str)
+    processSpeed = pyqtSignal(str)
+    processCurTime2 = pyqtSignal(str)
+    processSpeed2 = pyqtSignal(str)
+
+    def __init__(self):
+        super(HomePage, self).__init__()
+        self.layout = None
+        self.labels: list[WorkLabel] = []
+        self.worker: WorkerThread
+        self.btn_start: QPushButton
+        self.showMessageBox.connect(self.on_show_message_box_info)
+        self.labelDestroy.connect(self.on_label_destroy)
+        self.processCurTime.connect(self.on_process_cur_time_change)
+        self.processSpeed.connect(self.on_process_speed_change)
+        self.process_cur_time = "-"
+        self.process_speed = "-"
+        self.cur_clip_duration = 0
+        self.init_ui()
+
+    def init_ui(self):
+        layout = QVBoxLayout()
+        self.layout = QVBoxLayout()
+        layout.addLayout(self.layout)
+
+        btn_start = QPushButton(self)
+        btn_start.setText("开始")
+        btn_start.clicked.connect(self.handle_do)
+        btn_start.setGeometry(0, 560, 400, 40)
+        layout.addWidget(btn_start, 1, Qt.AlignBottom)
+        self.setLayout(layout)
+
+        self.btn_start = btn_start
+        self.resize(400, 600)
+        self.setWindowTitle("录播工作流")
+        self.setAcceptDrops(True)
+        self.show()
+
+    def add_label(self):
+        label1 = WorkLabel(self)
+        self.layout.addWidget(label1)
+        self.labels.append(label1)
+        return label1
+
+    def handle_do(self):
+        if len(self.labels) == 0:
+            return self.on_show_message_box_warn("提示", "请添加任务")
+        for label in self.labels:
+            if not label.have_job():
+                continue
+            _thread = WorkerThread(self, label)
+            _thread.start()
+            self.worker = _thread
+            _thread.started.connect(self.on_worker_start)
+            _thread.finished.connect(self.on_worker_stop)
+            break
+
+    def dragEnterEvent(self, a0: QtGui.QDragEnterEvent) -> None:
+        if a0.mimeData().hasUrls():
+            a0.accept()
+        else:
+            a0.ignore()
+
+    def dropEvent(self, a0: QtGui.QDropEvent) -> None:
+        label1 = self.add_label()
+        label1.dropEvent(a0)
+
+    def on_show_message_box_info(self, title: str, content: str):
+        QMessageBox.information(self, title, content, QMessageBox.Yes)
+
+    def on_show_message_box_warn(self, title: str, content: str):
+        QMessageBox.warning(self, title, content, QMessageBox.Yes)
+
+    def on_label_destroy(self, label: WorkLabel):
+        if label in self.labels:
+            self.labels.remove(label)
+
+    def on_worker_start(self):
+        self.btn_start.setDisabled(True)
+        self.btn_start.setText("正在处理")
+
+    def on_process_cur_time_change(self, s: str) -> None:
+        if self.process_cur_time == s:
+            return
+        self.process_cur_time = s
+        self.update_btn_process_text()
+
+    def on_process_speed_change(self, s: str) -> None:
+        if self.process_speed == s:
+            return
+        self.process_speed = s
+        self.update_btn_process_text()
+
+    def update_btn_process_text(self):
+        if self.process_cur_time != "" and self.process_speed != "":
+            self.btn_start.setText("{}@{}".format(self.process_cur_time, self.process_speed))
+        else:
+            self.btn_start.setText("Working")
+
+    def on_worker_stop(self):
+        self.btn_start.setDisabled(False)
+        self.btn_start.setText("开始")
+        self.worker = None
+        self.handle_do()
+
+
+class WorkerThread(QThread):
+    def __init__(self, app: HomePage, label: WorkLabel):
+        super(WorkerThread, self).__init__()
+        self.app = app
+        self.label = label
+
+    def get_video_real_duration(self, filename) -> float:
+        ffmpeg_process = subprocess.Popen([
+            FFMPEG_EXEC, "-hide_banner", "-progress", "-", "-loglevel", "error", "-i", filename, "-c", "copy", "-f",
+            "null", "-"
+        ], **subprocess_args(True))
+        _duration_str = self.handle_ffmpeg_output(ffmpeg_process.stdout)
+        ffmpeg_process.wait()
+        return duration_str_to_float(_duration_str)
+
+    def run(self) -> None:
+        try:
+            self.label.start_running()
+            job = self.label.get_job()
+            if job.type == Job.DANMAKU_ENCODE:
+                self.run_danmaku_encode(job)
+            elif job.type == Job.PURE_SPLIT:
+                self.quick_split_video(job.video)
+        except Exception as e:
+            print(e)
+            print(traceback.format_exc())
+        finally:
+            self.label.stop_running()
+
+    def run_danmaku_encode(self, job: Job):
+        base_danmaku = job.danmaku.pop(0)
+        time_shift = 0
+        try:
+            base_start_ts = get_file_start(base_danmaku)
+        except NoDanmakuException:
+            return
+        new_subtitle_name = danmaku_to_subtitle(base_danmaku, time_shift)
+        job.subtitles.append(new_subtitle_name)
+        for danmaku in job.danmaku:
+            try:
+                time_shift = diff_danmaku_files(base_danmaku, danmaku)
+            except NoDanmakuException:
+                continue
+            new_subtitle_name = danmaku_to_subtitle(danmaku, time_shift)
+            job.subtitles.append(new_subtitle_name)
+        # 压制
+        self.encode_video_with_subtitles(job.video, job.subtitles, base_start_ts)
+        for _f in job.subtitles:
+            if os.path.isfile(_f):
+                os.remove(_f)
+
+    def encode_video_with_subtitles(self, orig_filename: str, subtitles: list[str], base_ts: float):
+        new_filename = base_ts_to_filename(base_ts, False)
+        new_fullpath = os.path.join(VIDEO_OUTPUT_DIR, new_filename)
+        if FFMPEG_USE_NVIDIA_GPU:
+            process = get_encode_process_use_nvenc(orig_filename, subtitles, new_fullpath)
+        elif FFMPEG_USE_INTEL_GPU:
+            process = get_encode_process_use_intel(orig_filename, subtitles, new_fullpath)
+        else:
+            process = get_encode_process_use_cpu(orig_filename, subtitles, new_fullpath)
+        self.handle_ffmpeg_output(process.stdout)
+        process.wait()
+        return [new_fullpath]
+
+    def quick_split_video(self, file):
+        if not os.path.isfile(file):
+            raise FileNotFoundError(file)
+        file_name = os.path.split(file)[-1]
+        _create_dt = os.path.splitext(file_name)[0]
+        create_dt = datetime.strptime(_create_dt, "%Y%m%d_%H%M")
+        duration = self.get_video_real_duration(file)
+        current_sec = 0
+        if duration < VIDEO_CLIP_EACH_SEC:
+            print("[-]Less than each sec, skip")
+            return False
+        while current_sec < duration:
+            if (current_sec + VIDEO_CLIP_OVERFLOW_SEC * 2) > duration:
+                print("[-]Less than 2 overflow sec, skip")
+                break
+            current_dt = (create_dt + timedelta(seconds=current_sec)).strftime("%Y%m%d_%H%M_")
+            print("CUR_DT", current_dt)
+            print("BIAS_T", current_sec)
+            split_process = subprocess.Popen([
+                FFMPEG_EXEC, "-y", "-hide_banner", "-progress", "-", "-loglevel", "error",
+                "-ss", str(current_sec),
+                "-i", file, "-c:v", "copy", "-f", "mp4", "-c:a", "aac",
+                "-t", str(VIDEO_CLIP_EACH_SEC + VIDEO_CLIP_OVERFLOW_SEC),
+                "-fflags", "+genpts", "-shortest", "-movflags", "faststart",
+                os.path.join(VIDEO_OUTPUT_DIR, "{}.mp4".format(current_dt))
+            ], **subprocess_args(True))
+            self.handle_ffmpeg_output(split_process.stdout)
+            split_process.wait()
+            current_sec += VIDEO_CLIP_EACH_SEC
+        return True
+
+    def handle_ffmpeg_output(self, stdout: Optional[IO[bytes]]) -> str:
+        out_time = "0:0:0.0"
+        if stdout is None:
+            print("[!]STDOUT is null")
+            return out_time
+        speed = "0"
+        while True:
+            line = stdout.readline()
+            if line == b"":
+                break
+            if line.strip() == b"progress=end":
+                # 处理完毕
+                break
+            if line.startswith(b"out_time="):
+                out_time = line.replace(b"out_time=", b"").decode().strip()
+                self.app.processCurTime.emit(out_time)
+            if line.startswith(b"speed="):
+                speed = line.replace(b"speed=", b"").decode().strip()
+                self.app.processSpeed.emit(speed)
+        self.app.processSpeed.emit("")
+        self.app.processCurTime.emit("")
+        print("[ ]Speed:", out_time, "@", speed)
+        return out_time
+
+
+def duration_str_to_float(duration_str) -> float:
+    _duration = datetime.strptime(duration_str, "%H:%M:%S.%f") - datetime(1900, 1, 1)
+    return _duration.total_seconds()
+
+
+def base_ts_to_filename(start_ts: float, is_mp4=False) -> str:
+    base_start = datetime.fromtimestamp(start_ts)
+    if is_mp4:
+        return base_start.strftime("%Y%m%d_%H%M.mp4")
+    else:
+        return base_start.strftime("%Y%m%d_%H%M.flv")
+
+
+def danmaku_to_subtitle(file: Union[os.PathLike[str], str], time_shift: float):
+    new_subtitle_name = md5(file.encode("utf-8")).hexdigest() + ".ass"
+    process = subprocess.Popen((
+        DANMAKU_EXEC, "--ignore-warnings",
+        "-r", str(VIDEO_RESOLUTION), "-s", str(DANMAKU_SPEED), "-f", "5",
+        "-S", str(DANMAKU_FONT_SIZE), "-N", str(DANMAKU_FONT_NAME), "--showmsgbox", "FALSE",
+        "-O", "255", "-L", "1", "-D", "0",
+        "-o", "ass", new_subtitle_name, "-i", file, "-t", str(time_shift)
+    ), **subprocess_args(True))
+    process.wait()
+    return new_subtitle_name
+
+
+def get_encode_process_use_nvenc(orig_filename: str, subtitles: list[str], new_filename: str):
+    print("[+]Use Nvidia NvEnc Acceleration")
+    encode_process = subprocess.Popen([
+        FFMPEG_EXEC, *_common_ffmpeg_setting(),
+        "-hwaccel", "cuda", "-i", orig_filename, "-vf",
+        ",".join("subtitles=%s" % i for i in subtitles) + ",hwupload_cuda",
+        "-c:v", "h264_nvenc", "-preset:v", "p7", "-profile:v", "main",
+        "-tune", "hq", "-multipass", "2", "-cq", "30", "-rc-lookahead", "5", 
+        "-g", "600", "-keyint_min", "1", "-qdiff", "30", "-qcomp", "0.9",
+        *_common_ffmpeg_params(),
+        # "-t", "10",
+        new_filename
+    ], **subprocess_args(True))
+    return encode_process
+
+
+def get_encode_process_use_intel(orig_filename: str, subtitles: list[str], new_filename: str):
+    print("[+]Use Intel QSV Acceleration")
+    encode_process = subprocess.Popen([
+        FFMPEG_EXEC, *_common_ffmpeg_setting(),
+        "-hwaccel", "qsv", "-i", orig_filename, "-vf",
+        ",".join("subtitles=%s" % i for i in subtitles),
+        "-c:v", "h264_qsv", "-preset:v", "slow", "-profile:v", "main",
+        *_common_ffmpeg_params(),
+        # "-t", "10",
+        new_filename
+    ], **subprocess_args(True))
+    return encode_process
+
+
+def get_encode_process_use_cpu(orig_filename: str, subtitles: list[str], new_filename: str):
+    print("[+]Use CPU Encode")
+    encode_process = subprocess.Popen([
+        FFMPEG_EXEC, *_common_ffmpeg_setting(),
+        "-i", orig_filename, "-vf",
+        ",".join("subtitles=%s" % i for i in subtitles),
+        "-c:v", "h264", "-preset:v", "slow", "-profile:v", "main",
+        *_common_ffmpeg_params(),
+        # "-t", "10",
+        new_filename
+    ], **subprocess_args(True))
+    return encode_process
+
+
+# Create a set of arguments which make a ``subprocess.Popen`` (and
+# variants) call work with or without Pyinstaller, ``--noconsole`` or
+# not, on Windows and Linux. Typical use::
+#
+#   subprocess.call(['program_to_run', 'arg_1'], **subprocess_args())
+#
+# When calling ``check_output``::
+#
+#   subprocess.check_output(['program_to_run', 'arg_1'],
+#                           **subprocess_args(False))
+def subprocess_args(include_stdout=True):
+    # The following is true only on Windows.
+    if hasattr(subprocess, 'STARTUPINFO'):
+        # On Windows, subprocess calls will pop up a command window by default
+        # when run from Pyinstaller with the ``--noconsole`` option. Avoid this
+        # distraction.
+        si = subprocess.STARTUPINFO()
+        si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
+        # Windows doesn't search the path by default. Pass it an environment so
+        # it will.
+        env = os.environ
+    else:
+        si = None
+        env = None
+
+    # ``subprocess.check_output`` doesn't allow specifying ``stdout``::
+    #
+    #   Traceback (most recent call last):
+    #     File "test_subprocess.py", line 58, in <module>
+    #       **subprocess_args(stdout=None))
+    #     File "C:\Python27\lib\subprocess.py", line 567, in check_output
+    #       raise ValueError('stdout argument not allowed, it will be overridden.')
+    #   ValueError: stdout argument not allowed, it will be overridden.
+    #
+    # So, add it only if it's needed.
+    if include_stdout:
+        ret = {'stdout': subprocess.PIPE}
+    else:
+        ret = {}
+
+    # On Windows, running this from the binary produced by Pyinstaller
+    # with the ``--noconsole`` option requires redirecting everything
+    # (stdin, stdout, stderr) to avoid an OSError exception
+    # "[Error 6] the handle is invalid."
+    ret.update({'stdin': subprocess.PIPE,
+                'startupinfo': si,
+                'env': env})
+    return ret
+
+
+def check_exec(name: Union[os.PathLike[str], str]) -> bool:
+    if is_windows():
+        check_process = subprocess.Popen([
+            "where.exe", name
+        ], **subprocess_args(True))
+        check_process.wait()
+        return len(check_process.stdout.readlines()) > 0
+    elif is_linux():
+        check_process = subprocess.Popen([
+            "which", name
+        ])
+        check_process.wait()
+        return check_process.returncode == 0
+    else:
+        return False
+
+
+def is_windows() -> bool:
+    return platform.system().lower() == "windows"
+
+
+def is_linux() -> bool:
+    return platform.system().lower() == "linux"
+
+
+def check_all_prerequisite():
+    if not check_exec(DANMAKU_EXEC):
+        input("弹幕处理工具不存在")
+        exit(1)
+    if not check_exec(FFMPEG_EXEC):
+        input("FFMPEG工具不存在")
+        exit(1)
+
+
+def _common_ffmpeg_setting():
+    return (
+        "-y", "-hide_banner", "-progress", "-", "-loglevel", "error",
+    )
+
+
+def _common_ffmpeg_params():
+    return (
+        "-f", "mp4", "-c:a", "aac",
+        "-avoid_negative_ts", "1",
+        "-fflags", "+genpts", "-shortest"
+    )
+
+
+def main():
+    check_all_prerequisite()
+    app = QApplication(sys.argv)
+    page = HomePage()
+    sys.exit(app.exec_())
+
+
+if __name__ == '__main__':
+    load_config()
+    main()

danmaku_workflow_ffmpeg.spec

@@ -0,0 +1,40 @@
+# -*- mode: python ; coding: utf-8 -*-
+
+
+block_cipher = None
+
+
+a = Analysis(['danmaku_workflow_ffmpeg.py'],
+             pathex=[],
+             binaries=[],
+             datas=[],
+             hiddenimports=[],
+             hookspath=[],
+             hooksconfig={},
+             runtime_hooks=[],
+             excludes=[],
+             win_no_prefer_redirects=False,
+             win_private_assemblies=False,
+             cipher=block_cipher,
+             noarchive=False)
+pyz = PYZ(a.pure, a.zipped_data,
+             cipher=block_cipher)
+
+exe = EXE(pyz,
+          a.scripts,
+          a.binaries,
+          a.zipfiles,
+          a.datas,  
+          [],
+          name='录播工作流_ffmpeg',
+          debug=False,
+          bootloader_ignore_signals=False,
+          strip=False,
+          upx=True,
+          upx_exclude=[],
+          runtime_tmpdir=None,
+          console=False,
+          disable_windowed_traceback=False,
+          target_arch=None,
+          codesign_identity=None,
+          entitlements_file=None )