558 lines
19 KiB
Python
558 lines
19 KiB
Python
# 工作流
|
||
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),
|
||
"-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()
|