2022-12-05 10:40:27 +08:00

245 lines
9.0 KiB
Python

import os
import subprocess
from datetime import datetime, timedelta
from typing import IO
from config import FFMPEG_EXEC, FFMPEG_USE_HEVC, VIDEO_BITRATE, FFMPEG_USE_NVIDIA_GPU, VIDEO_CLIP_EACH_SEC, \
VIDEO_CLIP_OVERFLOW_SEC, \
FFMPEG_USE_INTEL_GPU, VIDEO_OUTPUT_DIR, VIDEO_CRF, VIDEO_GOP, FFMPEG_USE_VAAPI
from . import LOGGER
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 get_video_real_duration(filename):
ffmpeg_process = subprocess.Popen([
FFMPEG_EXEC, *_common_ffmpeg_setting(), "-i", filename, "-c", "copy", "-f", "null", "-"
], stdout=subprocess.PIPE)
return handle_ffmpeg_output(ffmpeg_process.stdout)
def encode_video_with_subtitles(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_HEVC:
if FFMPEG_USE_NVIDIA_GPU:
process = get_encode_hevc_process_use_nvenc(orig_filename, subtitles, new_fullpath)
elif FFMPEG_USE_VAAPI:
process = get_encode_hevc_process_use_vaapi(orig_filename, subtitles, new_fullpath)
elif FFMPEG_USE_INTEL_GPU:
process = get_encode_hevc_process_use_intel(orig_filename, subtitles, new_fullpath)
else:
process = get_encode_hevc_process_use_cpu(orig_filename, subtitles, new_fullpath)
else:
if FFMPEG_USE_NVIDIA_GPU:
process = get_encode_process_use_nvenc(orig_filename, subtitles, new_fullpath)
elif FFMPEG_USE_VAAPI:
process = get_encode_process_use_vaapi(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)
handle_ffmpeg_output(process.stdout)
process.wait()
return [new_fullpath]
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(),
"-i", orig_filename, "-vf",
",".join("subtitles=%s" % i for i in subtitles) + ",hwupload_cuda",
"-c:v", "h264_nvenc",
*_common_ffmpeg_params(),
# "-t", "10",
new_filename
], stdout=subprocess.PIPE)
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",
*_common_ffmpeg_params(),
# "-t", "10",
new_filename
], stdout=subprocess.PIPE)
return encode_process
def get_encode_process_use_vaapi(orig_filename: str, subtitles: list[str], new_filename: str):
print("[+]Use VAAPI Acceleration")
encode_process = subprocess.Popen([
FFMPEG_EXEC, *_common_ffmpeg_setting(),
"-hwaccel", "vaapi", "-hwaccel_output_format", "vaapi", "-i", orig_filename, "-vf",
"hwmap=mode=read+write+direct,format=nv12," +
"".join("subtitles=%s," % i for i in subtitles) + "hwmap",
"-c:v", "h264_vaapi",
*_common_ffmpeg_params(),
# "-t", "10",
new_filename
], stdout=subprocess.PIPE)
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",
*_common_ffmpeg_params(),
# "-t", "10",
new_filename
], stdout=subprocess.PIPE)
return encode_process
def get_encode_hevc_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(),
"-i", orig_filename, "-vf",
"".join("subtitles=%s," % i for i in subtitles) + ",hwupload_cuda",
"-c:v", "hevc_nvenc",
*_common_ffmpeg_params(),
# "-t", "10",
new_filename
], stdout=subprocess.PIPE)
return encode_process
def get_encode_hevc_process_use_vaapi(orig_filename: str, subtitles: list[str], new_filename: str):
print("[+]Use VAAPI Acceleration")
encode_process = subprocess.Popen([
FFMPEG_EXEC, *_common_ffmpeg_setting(),
"-hwaccel", "vaapi", "-hwaccel_output_format", "vaapi", "-i", orig_filename, "-vf",
"hwmap=mode=read+write+direct,format=nv12," +
"".join("subtitles=%s," % i for i in subtitles) + "hwmap",
"-c:v", "hevc_vaapi",
*_common_ffmpeg_params(),
# "-t", "10",
new_filename
], stdout=subprocess.PIPE)
return encode_process
def get_encode_hevc_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", "hevc_qsv",
*_common_ffmpeg_params(),
# "-t", "10",
new_filename
], stdout=subprocess.PIPE)
return encode_process
def get_encode_hevc_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", "hevc",
*_common_ffmpeg_params(),
# "-t", "10",
new_filename
], stdout=subprocess.PIPE)
return encode_process
def handle_ffmpeg_output(stdout: IO[bytes]) -> str:
out_time = "0:0:0.0"
speed = "0"
if stdout is None:
print("[!]STDOUT is null")
return out_time
_i = 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()
if line.startswith(b"speed="):
speed = line.replace(b"speed=", b"").decode().strip()
_i += 1
if _i % 300 == 150:
LOGGER.debug("[>]Speed:{}@{}".format(out_time, speed))
LOGGER.debug("[ ]Speed:{}@{}".format(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 quick_split_video(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_str = get_video_real_duration(file)
duration = duration_str_to_float(_duration_str)
current_sec = 0
_video_parts = []
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, *_common_ffmpeg_setting(),
"-ss", str(current_sec),
"-i", file, "-c:v", "copy", "-c:a", "aac", "-f", "mp4",
"-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))
], stdout=subprocess.PIPE)
handle_ffmpeg_output(split_process.stdout)
split_process.wait()
_video_parts.append({
"base_path": VIDEO_OUTPUT_DIR,
"file": "{}.mp4".format(current_dt),
})
current_sec += VIDEO_CLIP_EACH_SEC
return _video_parts
def _common_ffmpeg_setting():
return (
"-y", "-hide_banner", "-progress", "-", "-loglevel", "error",
)
def _common_ffmpeg_params():
return (
"-f", "mp4", "-b:v", VIDEO_BITRATE, "-c:a", "aac",
"-preset:v", "fast", "-profile:v", "main", "-avoid_negative_ts", "1",
"-qmin", "18", "-qmax", "38", "-crf", str(VIDEO_CRF), "-g:v", str(VIDEO_GOP),
"-fflags", "+genpts", "-shortest"
)