# 工作流 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, DANMAKU_OPACITY, \ 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), "-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 # **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()