426 lines
14 KiB
Python
426 lines
14 KiB
Python
# 工作流
|
||
import datetime
|
||
import os.path
|
||
import subprocess
|
||
import sys
|
||
from typing import Optional, IO, AnyStr
|
||
|
||
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
|
||
|
||
DANMAKU_PATH = "workspace"
|
||
DANMAKU_FACTORY_CLI = "DFCLI.exe"
|
||
FFMPEG_TOOL = "ffmpeg.exe"
|
||
SPLIT_PATH = "split"
|
||
SPLIT_TOOL = "quick_split_video.exe"
|
||
HOME_PATH = os.path.realpath(".")
|
||
|
||
|
||
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.workVideo is None or self.finished:
|
||
self.parent().labelDestroy.emit(self)
|
||
return self.deleteLater()
|
||
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)
|
||
|
||
def __init__(self):
|
||
super(HomePage, self).__init__()
|
||
self.layout = None
|
||
self.labels: list[WorkLabel] = []
|
||
self.worker: Optional[WorkerThread] = None
|
||
self.btn_start: Optional[QPushButton] = None
|
||
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.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:
|
||
self.process_cur_time = s
|
||
self.update_btn_process_text()
|
||
|
||
def on_process_speed_change(self, s: str)->None:
|
||
self.process_speed = s
|
||
self.update_btn_process_text()
|
||
|
||
def update_btn_process_text(self):
|
||
self.btn_start.setText("{}@{}".format(self.process_cur_time, self.process_speed))
|
||
|
||
def on_worker_stop(self):
|
||
self.btn_start.setDisabled(False)
|
||
self.btn_start.setText("开始")
|
||
self.worker = None
|
||
|
||
|
||
class WorkerThread(QThread):
|
||
def __init__(self, app: HomePage, label: WorkLabel):
|
||
super(WorkerThread, self).__init__()
|
||
self.app = app
|
||
self.label = label
|
||
|
||
def run(self) -> None:
|
||
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.run_pure_split(job)
|
||
self.label.stop_running()
|
||
|
||
def run_danmaku_encode(self, job: Job):
|
||
os.chdir(DANMAKU_PATH)
|
||
base_danmaku = job.danmaku.pop(0)
|
||
time_shift = 0
|
||
base_start_ts = get_file_start(base_danmaku)
|
||
new_subtitle_name = os.path.basename(base_danmaku) + ".ass"
|
||
subprocess.Popen((
|
||
DANMAKU_FACTORY_CLI,
|
||
"-r", "1280x720", "-s", "10", "-f", "5",
|
||
"-S", "40", "-N", "\"Sarasa Term SC\"", "--showmsgbox", "FALSE",
|
||
"-O", "255", "-L", "1", "-D", "0",
|
||
"-o", "ass", new_subtitle_name, "-i", base_danmaku, "-t", str(time_shift)
|
||
), **subprocess_args()).wait()
|
||
print(new_subtitle_name)
|
||
job.subtitles.append(new_subtitle_name)
|
||
for danmaku in job.danmaku:
|
||
time_shift = diff_danmaku_files(base_danmaku, danmaku)
|
||
new_subtitle_name = os.path.basename(danmaku) + ".ass"
|
||
subprocess.Popen((
|
||
DANMAKU_FACTORY_CLI,
|
||
"-r", "1280x720", "-s", "10", "-f", "5",
|
||
"-S", "40", "-N", "\"Sarasa Term SC\"", "--showmsgbox", "FALSE",
|
||
"-O", "255", "-L", "1", "-D", "0",
|
||
"-o", "ass", new_subtitle_name, "-i", danmaku, "-t", str(time_shift)
|
||
), **subprocess_args()).wait()
|
||
print(new_subtitle_name)
|
||
job.subtitles.append(new_subtitle_name)
|
||
# 压制
|
||
base_start = datetime.datetime.fromtimestamp(base_start_ts)
|
||
new_file_name = base_start.strftime("%Y%m%d_%H%M.flv")
|
||
encode_process = subprocess.Popen([
|
||
FFMPEG_TOOL, "-hide_banner", "-progress", "-", "-v", "0", "-y",
|
||
"-i", job.video, "-vf", ",".join("subtitles=%s" % i for i in job.subtitles) + ",hwupload_cuda",
|
||
"-c:a", "copy", "-c:v", "h264_nvenc", "-f", "mp4",
|
||
"-preset:v", "fast", "-profile:v", "high", "-level", "4.1",
|
||
"-b:v", "2.5M", "-bufsize:v", "5M", "-rc:v", "vbr_hq", "-bf:v", "3",
|
||
"-qmin", "10", "-qmax", "52", "-crf", "16",
|
||
# "-t", "10",
|
||
os.path.join(HOME_PATH, SPLIT_PATH, new_file_name)
|
||
], **subprocess_args(True))
|
||
self.handle_ffmpeg_output(encode_process.stdout)
|
||
encode_process.wait()
|
||
for _f in job.subtitles:
|
||
if os.path.isfile(_f):
|
||
os.remove(_f)
|
||
os.chdir(HOME_PATH)
|
||
os.chdir(SPLIT_PATH)
|
||
split_process = subprocess.Popen([
|
||
SPLIT_TOOL, new_file_name
|
||
], **subprocess_args(True))
|
||
self.handle_ffmpeg_output(split_process.stdout)
|
||
split_process.wait()
|
||
os.chdir(HOME_PATH)
|
||
|
||
def run_pure_split(self, job: Job):
|
||
os.chdir(SPLIT_PATH)
|
||
split_process = subprocess.Popen([
|
||
SPLIT_TOOL, job.video
|
||
], **subprocess_args(True))
|
||
self.handle_ffmpeg_output(split_process.stdout)
|
||
split_process.wait()
|
||
os.chdir(HOME_PATH)
|
||
|
||
def handle_ffmpeg_output(self, stderr: IO[bytes]) -> None:
|
||
while True:
|
||
line = stderr.readline()
|
||
if line == b"":
|
||
break
|
||
if line.startswith(b"out_time="):
|
||
cur_time = line.replace(b"out_time=", b"").decode()
|
||
self.app.processCurTime.emit(cur_time.strip())
|
||
if line.startswith(b"speed="):
|
||
speed = line.replace(b"speed=", b"").decode()
|
||
self.app.processSpeed.emit(speed.strip())
|
||
|
||
|
||
# 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,
|
||
'stderr': subprocess.PIPE,
|
||
'startupinfo': si,
|
||
'env': env})
|
||
return ret
|
||
|
||
|
||
def check_all_prerequisite():
|
||
if not os.path.isdir(DANMAKU_PATH):
|
||
os.mkdir(DANMAKU_PATH)
|
||
if not os.path.isdir(SPLIT_PATH):
|
||
os.mkdir(SPLIT_PATH)
|
||
os.chdir(DANMAKU_PATH)
|
||
validate_process = subprocess.Popen([
|
||
"where.exe", DANMAKU_FACTORY_CLI
|
||
], **subprocess_args(True))
|
||
if len(validate_process.stdout.readlines()) == 0:
|
||
input("弹幕处理工具不存在")
|
||
exit(1)
|
||
os.chdir(HOME_PATH)
|
||
os.chdir(SPLIT_PATH)
|
||
validate_process = subprocess.Popen([
|
||
"where.exe", SPLIT_TOOL
|
||
], **subprocess_args(True))
|
||
if len(validate_process.stdout.readlines()) == 0:
|
||
input("视频分割工具不存在")
|
||
exit(1)
|
||
os.chdir(HOME_PATH)
|
||
validate_process = subprocess.Popen([
|
||
"where.exe", FFMPEG_TOOL
|
||
], **subprocess_args(True))
|
||
if len(validate_process.stdout.readlines()) == 0:
|
||
input("FFMPEG工具不存在")
|
||
exit(1)
|
||
os.chdir(HOME_PATH)
|
||
|
||
|
||
def main():
|
||
check_all_prerequisite()
|
||
app = QApplication(sys.argv)
|
||
page = HomePage()
|
||
sys.exit(app.exec_())
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|