766 lines
32 KiB
Python
766 lines
32 KiB
Python
import base64
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
import subprocess
|
|
import os
|
|
import platform
|
|
import tarfile
|
|
import zipfile
|
|
from functools import lru_cache
|
|
import shutil
|
|
from PIL import Image, PngImagePlugin
|
|
import json
|
|
from tqdm import tqdm
|
|
import io
|
|
from pathlib import Path
|
|
import ssl
|
|
import mimetypes
|
|
import imghdr
|
|
import piexif
|
|
from .public import args, find_project_root, generate_md5_uid_timestamp, determine_file_type, find_project_custiom_nodes_path
|
|
import folder_paths
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
input_directory = folder_paths.get_input_directory()
|
|
os.makedirs(input_directory, exist_ok=True)
|
|
save_input_directory = input_directory + '/temp'
|
|
os.makedirs(save_input_directory, exist_ok=True)
|
|
FFMPEG_URLS = {
|
|
"amd64-static": "https://tt-1254127940.cos.ap-guangzhou.myqcloud.com/ffmpeg/ffmpeg-release-amd64-static.tar.xz",
|
|
"i686-static": "https://tt-1254127940.cos.ap-guangzhou.myqcloud.com/ffmpeg/ffmpeg-release-i686-static.tar.xz",
|
|
"arm64-static": "https://tt-1254127940.cos.ap-guangzhou.myqcloud.com/ffmpeg/ffmpeg-release-arm64-static.tar.xz",
|
|
"armhf-static": "https://tt-1254127940.cos.ap-guangzhou.myqcloud.com/ffmpeg/ffmpeg-release-armhf-static.tar.xz",
|
|
"armel-static": "https://tt-1254127940.cos.ap-guangzhou.myqcloud.com/ffmpeg/ffmpeg-release-armel-static.tar.xz",
|
|
"osx64": "https://tt-1254127940.cos.ap-guangzhou.myqcloud.com/ffmpeg/ffmpeg-116599-g43cde54fc1.zip",
|
|
"win64-full": "https://tt-1254127940.cos.ap-guangzhou.myqcloud.com/ffmpeg/ffmpeg-release-essentials.zip"
|
|
}
|
|
ssl._create_default_https_context = ssl._create_unverified_context
|
|
ffmpeg_path_exe = ''
|
|
ffprobe_exe_path = ''
|
|
temp_path = find_project_custiom_nodes_path() + 'ComfyUI_Bxb/temp_bxb/'
|
|
try:
|
|
resample_filter = Image.Resampling.LANCZOS
|
|
except AttributeError:
|
|
resample_filter = Image.LANCZOS
|
|
def file_to_base64(file_path):
|
|
mime_type, _ = mimetypes.guess_type(file_path)
|
|
if mime_type is None:
|
|
raise ValueError("无法识别的文件类型")
|
|
with open(file_path, "rb") as file:
|
|
base64_encoded = base64.b64encode(file.read()).decode('utf-8')
|
|
base64_prefix = f"data:{mime_type};base64,"
|
|
return base64_prefix + base64_encoded
|
|
def compress_image(input_image_path, output_image_path=None, target_width=None, quality=85):
|
|
|
|
img = Image.open(input_image_path)
|
|
if target_width is None:
|
|
target_width = img.size[0]
|
|
width_percent = target_width / float(img.size[0])
|
|
target_height = int((float(img.size[1]) * float(width_percent)))
|
|
img_resized = img.resize((target_width, target_height), Image.LANCZOS)
|
|
if output_image_path is None:
|
|
output_image_path = input_image_path
|
|
if output_image_path == input_image_path:
|
|
os.remove(input_image_path)
|
|
img_resized.save(output_image_path, quality=quality, optimize=True)
|
|
def cut_and_compress_video(input_video_path, output_video_path=None, target_width=None, duration=None):
|
|
|
|
probe_command = [
|
|
ffmpeg_path_exe, '-v', 'error', '-select_streams', 'v:0', '-show_entries',
|
|
'stream=width,height', '-of', 'csv=p=0:s=x', input_video_path
|
|
]
|
|
result = subprocess.run(probe_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
original_width, original_height = map(int, result.stdout.strip().split('x'))
|
|
if target_width is None:
|
|
target_width = original_width
|
|
width_ratio = target_width / float(original_width)
|
|
target_height = int(original_height * width_ratio)
|
|
if output_video_path is None:
|
|
output_video_path = input_video_path
|
|
command = [
|
|
ffmpeg_path_exe,
|
|
'-i', input_video_path,
|
|
'-vf', f'scale={target_width}:{target_height}',
|
|
'-r', '24',
|
|
'-c:v', 'libx264',
|
|
'-preset', 'medium',
|
|
'-crf', '28',
|
|
'-c:a', 'aac',
|
|
'-b:a', '128k',
|
|
'-strict', 'experimental',
|
|
'-y',
|
|
output_video_path
|
|
]
|
|
if duration is not None:
|
|
command.insert(1, '-t')
|
|
command.insert(2, str(duration))
|
|
subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
|
|
def get_platform():
|
|
architecture = platform.machine()
|
|
if sys.platform.startswith("linux"):
|
|
if architecture == "x86_64":
|
|
return "amd64-static"
|
|
elif architecture == "i686":
|
|
return "i686-static"
|
|
elif architecture == "aarch64":
|
|
return "arm64-static"
|
|
elif architecture == "armv7l":
|
|
return "armhf-static"
|
|
elif architecture == "armel":
|
|
return "armel-static"
|
|
elif sys.platform.startswith("darwin"):
|
|
return "osx64"
|
|
elif sys.platform.startswith("win"):
|
|
return "win64-full"
|
|
else:
|
|
raise RuntimeError("Unsupported platform")
|
|
def download_ffmpeg(target_dir):
|
|
def reporthook(block_num, block_size, total_size):
|
|
if reporthook.tbar is None:
|
|
reporthook.tbar = tqdm(total=total_size, unit='B', unit_scale=True)
|
|
downloaded = block_num * block_size
|
|
if downloaded < total_size:
|
|
reporthook.tbar.update(block_size)
|
|
else:
|
|
reporthook.tbar.close()
|
|
reporthook.tbar = None
|
|
reporthook.tbar = None
|
|
plat = get_platform()
|
|
url = FFMPEG_URLS.get(plat)
|
|
if not url:
|
|
raise RuntimeError(f"不支持的平台或配置: {plat}")
|
|
target_dir = Path(target_dir).resolve()
|
|
target_dir.mkdir(parents=True, exist_ok=True)
|
|
file_extension = url.split('.')[-1]
|
|
archive_path = target_dir / f"ffmpeg-{plat}.{file_extension}"
|
|
print(f"正在下载 FFmpeg for {plat}...")
|
|
no_proxy_handler = urllib.request.ProxyHandler({})
|
|
opener = urllib.request.build_opener(no_proxy_handler)
|
|
with opener.open(url) as response:
|
|
with open(archive_path, 'wb') as out_file:
|
|
block_size = 8192
|
|
total_size = int(response.headers.get('Content-Length', 0))
|
|
reporthook(0, block_size, total_size)
|
|
while True:
|
|
buffer = response.read(block_size)
|
|
if not buffer:
|
|
break
|
|
out_file.write(buffer)
|
|
reporthook(len(buffer) // block_size, block_size, total_size)
|
|
print(f"正在解压 FFmpeg...")
|
|
extracted_dir = target_dir / "ffmpeg"
|
|
if file_extension == "zip":
|
|
with zipfile.ZipFile(archive_path, 'r') as zip_ref:
|
|
zip_ref.extractall(extracted_dir)
|
|
elif file_extension == "xz":
|
|
with tarfile.open(archive_path, 'r:xz') as tar_ref:
|
|
tar_ref.extractall(extracted_dir)
|
|
ffmpeg_exe = find_ffmpeg_executable(extracted_dir)
|
|
os.remove(archive_path)
|
|
if not ffmpeg_exe or not ffmpeg_exe.exists():
|
|
raise RuntimeError(f"从 {archive_path} 提取 FFmpeg 可执行文件失败")
|
|
return str(ffmpeg_exe)
|
|
@lru_cache()
|
|
def find_ffmpeg_executable(search_dir):
|
|
|
|
if not os.path.exists(search_dir):
|
|
return None
|
|
os.chmod(search_dir, 0o755)
|
|
for root, dirs, files in os.walk(search_dir):
|
|
for file in files:
|
|
if file.lower() == "ffmpeg.exe" or file.lower() == "ffmpeg":
|
|
return Path(root) / file
|
|
return None
|
|
@lru_cache()
|
|
def find_ffprobe_executable(search_dir):
|
|
|
|
for root, dirs, files in os.walk(search_dir):
|
|
for file in files:
|
|
if file.lower() == "ffprobe.exe" or file.lower() == "ffprobe":
|
|
return Path(root) / file
|
|
return None
|
|
@lru_cache()
|
|
def get_ffmpeg_executable():
|
|
global ffmpeg_path_exe, ffprobe_exe_path
|
|
target_dir = find_project_custiom_nodes_path() + 'ComfyUI_Bxb/tools/'
|
|
ffmpeg_exe = find_ffmpeg_executable(target_dir + "/ffmpeg")
|
|
if ffmpeg_exe and ffmpeg_exe.exists() and is_valid_exe(ffmpeg_exe):
|
|
ffmpeg_path_exe = str(ffmpeg_exe)
|
|
ffprobe_exe = find_ffprobe_executable(target_dir + "/ffmpeg")
|
|
ffprobe_exe_path = str(ffprobe_exe)
|
|
os.chmod(ffprobe_exe_path, 0o755)
|
|
return str(ffmpeg_exe)
|
|
remove_target_dir = target_dir + "/ffmpeg"
|
|
if os.path.exists(remove_target_dir):
|
|
shutil.rmtree(remove_target_dir)
|
|
exe = download_ffmpeg(target_dir)
|
|
if exe and os.path.isfile(exe) and is_valid_exe(exe):
|
|
ffmpeg_path_exe = exe
|
|
ffprobe_exe = find_ffprobe_executable(target_dir + "/ffmpeg")
|
|
ffprobe_exe_path = str(ffprobe_exe)
|
|
os.chmod(ffprobe_exe_path, 0o755)
|
|
return exe
|
|
raise RuntimeError("FFmpeg not found or is invalid.")
|
|
def is_valid_exe(exe):
|
|
|
|
try:
|
|
os.chmod(exe, 0o755)
|
|
result = subprocess.run([exe, '-version'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
return result.returncode == 0
|
|
except Exception:
|
|
return False
|
|
from io import BytesIO
|
|
def get_image_dimensions(input_file, custom_data=None):
|
|
with Image.open(input_file) as img:
|
|
width, height = img.size
|
|
img_format = img.format
|
|
if not custom_data or (img_format != 'PNG' and img_format != 'JPEG'):
|
|
file_size_bytes = os.path.getsize(input_file)
|
|
file_size_mb = file_size_bytes / (1024 * 1024)
|
|
return width, height, file_size_mb
|
|
img_byte_arr = BytesIO()
|
|
if img_format == 'PNG':
|
|
meta = PngImagePlugin.PngInfo()
|
|
for key, value in custom_data.items():
|
|
if isinstance(value, bytes):
|
|
value = base64.b64encode(value).decode('utf-8')
|
|
meta.add_text(key, value)
|
|
img.save(img_byte_arr, format="PNG", pnginfo=meta)
|
|
elif img_format == 'JPEG':
|
|
img = img.convert("RGB")
|
|
exif_dict = {"0th": {}, "Exif": {}, "GPS": {}, "1st": {}, "thumbnail": None}
|
|
for key, value in custom_data.items():
|
|
exif_dict["0th"][piexif.ImageIFD.Make] = str(value)
|
|
exif_bytes = piexif.dump(exif_dict)
|
|
img.save(img_byte_arr, format="JPEG", exif=exif_bytes)
|
|
img_byte_arr.seek(0)
|
|
image_bytes = img_byte_arr.getvalue()
|
|
with open(input_file, "wb") as f:
|
|
f.write(image_bytes)
|
|
file_size_bytes = len(image_bytes)
|
|
file_size_mb = file_size_bytes / (1024 * 1024)
|
|
return width, height, file_size_mb
|
|
def get_video_dimensions(input_file):
|
|
command = [
|
|
ffprobe_exe_path,
|
|
'-v', 'error',
|
|
'-select_streams', 'v:0',
|
|
'-show_entries', 'stream=width,height',
|
|
'-of', 'json',
|
|
input_file
|
|
]
|
|
result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
result_json = json.loads(result.stdout)
|
|
width = result_json['streams'][0]['width']
|
|
height = result_json['streams'][0]['height']
|
|
file_size_bytes = os.path.getsize(input_file)
|
|
file_size_mb = file_size_bytes / (1024 * 1024)
|
|
return width, height, file_size_mb
|
|
def cut_video(input_file, start_seconds, end_seconds, output_file, width, height, fps=24, threads=4):
|
|
|
|
duration = end_seconds - start_seconds
|
|
if duration <= 0:
|
|
raise ValueError("结束时间必须大于开始时间")
|
|
scale_filter = f'scale={width}:{height}:force_original_aspect_ratio=increase'
|
|
crop_filter = f'crop={width}:{height}'
|
|
filter_complex = f'{scale_filter},{crop_filter}'
|
|
command = [
|
|
ffmpeg_path_exe,
|
|
'-ss', str(start_seconds),
|
|
'-i', input_file,
|
|
'-t', str(duration),
|
|
'-r', str(fps),
|
|
'-vf', filter_complex,
|
|
'-c:v', 'libx264',
|
|
'-c:a', 'aac',
|
|
'-threads', str(threads),
|
|
'-strict', 'experimental',
|
|
'-y',
|
|
output_file
|
|
]
|
|
subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
|
|
return duration
|
|
def loop_video_to_duration(input_file, output_file, target_duration):
|
|
|
|
command = [
|
|
ffmpeg_path_exe,
|
|
'-stream_loop', '-1',
|
|
'-i', input_file,
|
|
'-t', str(target_duration),
|
|
'-c:v', 'libx264',
|
|
'-c:a', 'aac',
|
|
'-threads', '8',
|
|
'-y',
|
|
output_file
|
|
]
|
|
subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
|
|
def merge_videos_horizontally(input_file1, input_file2, output_file='temp_frames'):
|
|
|
|
command = [
|
|
ffmpeg_path_exe,
|
|
'-i', input_file1,
|
|
'-i', input_file2,
|
|
'-filter_complex', '[0:v][1:v]hstack=inputs=2[v]',
|
|
'-map', '[v]',
|
|
'-c:v', 'libx264',
|
|
'-threads', '8',
|
|
'-y',
|
|
output_file
|
|
]
|
|
subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
|
|
def get_file_count_in_directory():
|
|
directory = Path(temp_path)
|
|
if not directory.exists():
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
file_count = len([file for file in directory.iterdir() if file.is_file()])
|
|
return file_count
|
|
def resize_and_crop_image(image_path, output_path, width, height):
|
|
|
|
command = [
|
|
ffmpeg_path_exe,
|
|
'-i', str(image_path),
|
|
'-vf', f"scale={width}:{height}:force_original_aspect_ratio=increase,crop={width}:{height}",
|
|
'-q:v', '2',
|
|
str(output_path)
|
|
]
|
|
try:
|
|
subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Error during resizing and cropping image: {e}")
|
|
def process_image(input_image_path_left=None, input_image_path_right=None, output_image_path="output.png",
|
|
canvas_size=(800, 600), overlay_image_path=None):
|
|
canvas = Image.new('RGBA', canvas_size, (0, 0, 0, 0))
|
|
half_width = canvas_size[0] // 2
|
|
target_height = canvas_size[1]
|
|
if input_image_path_left:
|
|
img_left = Image.open(input_image_path_left).convert("RGBA")
|
|
img_width, img_height = img_left.size
|
|
scale = max(half_width / img_width, target_height / img_height)
|
|
new_size = (int(img_width * scale), int(img_height * scale))
|
|
img_resized = img_left.resize(new_size, resample_filter)
|
|
img_cropped = img_resized.crop((
|
|
(img_resized.width - half_width) // 2,
|
|
(img_resized.height - target_height) // 2,
|
|
(img_resized.width + half_width) // 2,
|
|
(img_resized.height + target_height) // 2
|
|
))
|
|
canvas.paste(img_cropped, (0, 0), img_cropped)
|
|
if input_image_path_right:
|
|
img_right = Image.open(input_image_path_right).convert("RGBA")
|
|
img_width, img_height = img_right.size
|
|
scale = max(half_width / img_width, target_height / img_height)
|
|
new_size = (int(img_width * scale), int(img_height * scale))
|
|
img_resized = img_right.resize(new_size, resample_filter)
|
|
img_cropped = img_resized.crop((
|
|
(img_resized.width - half_width) // 2,
|
|
(img_resized.height - target_height) // 2,
|
|
(img_resized.width + half_width) // 2,
|
|
(img_resized.height + target_height) // 2
|
|
))
|
|
canvas.paste(img_cropped, (half_width, 0), img_cropped)
|
|
if overlay_image_path and os.path.exists(overlay_image_path):
|
|
overlay_img = Image.open(overlay_image_path).convert("RGBA")
|
|
overlay_scale = min(canvas_size[0] / 2 / overlay_img.width, canvas_size[1] / 2 / overlay_img.height)
|
|
overlay_new_size = (int(overlay_img.width * overlay_scale), int(overlay_img.height * overlay_scale))
|
|
if input_image_path_left is not None and input_image_path_right is not None and os.path.exists(input_image_path_left) and os.path.exists(input_image_path_right):
|
|
overlay_new_size = (int(canvas_size[0] / 15), int(canvas_size[0] / 15))
|
|
overlay_img_resized = overlay_img.resize(overlay_new_size, resample_filter)
|
|
overlay_x = (canvas_size[0] - overlay_new_size[0]) // 2
|
|
overlay_y = (canvas_size[1] - overlay_new_size[1]) // 2
|
|
canvas.paste(overlay_img_resized, (overlay_x, overlay_y), overlay_img_resized)
|
|
canvas.save(output_image_path, format='PNG', optimize=True)
|
|
def resize_and_crop(image, target_width, target_height):
|
|
|
|
scale = max(target_width / image.width, target_height / image.height)
|
|
new_size = (int(image.width * scale), int(image.height * scale))
|
|
resized_image = image.resize(new_size, resample_filter)
|
|
left = (resized_image.width - target_width) // 2
|
|
top = (resized_image.height - target_height) // 2
|
|
right = (resized_image.width + target_width) // 2
|
|
bottom = (resized_image.height + target_height) // 2
|
|
cropped_image = resized_image.crop((left, top, right, bottom))
|
|
return cropped_image
|
|
def process_and_merge_image_video(image_file, video_file, output_file, overlay_image, side='left', target_width=330, target_height=480,
|
|
start_seconds=0, end_seconds=None, fps=24):
|
|
|
|
temp_black_background = ''
|
|
if not image_file:
|
|
img = Image.new('RGBA', (target_width, target_height), (0, 0, 0, 255))
|
|
temp_black_background = temp_path + generate_md5_uid_timestamp(
|
|
'temp_black_background') + 'temp_black_background.png'
|
|
img.save(temp_black_background)
|
|
image_file = temp_black_background
|
|
image = Image.open(image_file).convert("RGBA")
|
|
image_resized = resize_and_crop(image, target_width, target_height)
|
|
video_resized_png = temp_path + generate_md5_uid_timestamp(str(time.process_time())) + 'resized_video.png'
|
|
image_resized.save(video_resized_png)
|
|
video_file_temp = ''
|
|
if end_seconds is not None:
|
|
video_file_temp = temp_path + generate_md5_uid_timestamp(str(time.process_time())) + 'cut_video.mp4'
|
|
cut_video(video_file, start_seconds, end_seconds, video_file_temp, target_width, target_height, fps)
|
|
video_file = video_file_temp
|
|
scale_filter = f'scale={target_width}:{target_height}:force_original_aspect_ratio=increase'
|
|
crop_filter = f'crop={target_width}:{target_height}'
|
|
filter_complex = f'{scale_filter},{crop_filter}'
|
|
overlay_new_width = int(target_width * 2 / 15)
|
|
overlay_new_height = int(target_width * 2 / 15)
|
|
overlay_x = (target_width * 2 - overlay_new_width) // 2
|
|
overlay_y = (target_height - overlay_new_height) // 2
|
|
if side == 'left':
|
|
filter_complex = (
|
|
f"[0:v]{scale_filter},{crop_filter}[vid];"
|
|
f"[1:v]scale={target_width}:{target_height}[img];"
|
|
f"[img][vid]hstack=inputs=2[base];"
|
|
f"[2:v]scale={overlay_new_width}:{overlay_new_height}[ovrl];"
|
|
f"[base][ovrl]overlay={overlay_x}:{overlay_y}[v]"
|
|
)
|
|
else:
|
|
filter_complex = (
|
|
f"[0:v]{scale_filter},{crop_filter}[vid];"
|
|
f"[1:v]scale={target_width}:{target_height}[img];"
|
|
f"[vid][img]hstack=inputs=2[base];"
|
|
f"[2:v]scale={overlay_new_width}:{overlay_new_height}[ovrl];"
|
|
f"[base][ovrl]overlay={overlay_x}:{overlay_y}[v]"
|
|
)
|
|
command = [
|
|
ffmpeg_path_exe,
|
|
'-i', video_file,
|
|
'-i', video_resized_png,
|
|
'-i', 'pipe:0',
|
|
'-filter_complex', filter_complex,
|
|
'-map', '[v]',
|
|
'-c:v', 'libx264',
|
|
'-pix_fmt', 'yuv420p',
|
|
'-threads', '4',
|
|
'-an',
|
|
'-y',
|
|
output_file
|
|
]
|
|
with open(overlay_image, 'rb') as overlay_img_file:
|
|
subprocess.run(command, stdin=overlay_img_file, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
|
|
os.remove(video_resized_png)
|
|
if temp_black_background:
|
|
os.remove(temp_black_background)
|
|
if video_file_temp:
|
|
os.remove(video_file_temp)
|
|
def process_and_merge_videos(input_file1, input_file2, start_seconds1, end_seconds1, start_seconds2, end_seconds2,
|
|
output_file, overlay_image, fps=24):
|
|
global CANVAS_MAX_HEIGHT, CANVAS_MAX_WIDTH
|
|
video_width1, video_height1, video_size = get_video_dimensions(input_file2)
|
|
get_file_count_in_directory()
|
|
cut1 = temp_path + generate_md5_uid_timestamp('cut1.mp4') + 'cut1.mp4'
|
|
cut2 = temp_path + generate_md5_uid_timestamp('cut1.mp4') + 'cut2.mp4'
|
|
cut1_looped = temp_path + generate_md5_uid_timestamp('cut1.mp4') + 'cut1_looped.mp4'
|
|
cut2_looped = temp_path + generate_md5_uid_timestamp('cut1.mp4') + 'cut2_looped.mp4'
|
|
max_width = CANVAS_MAX_WIDTH / 2
|
|
max_height = CANVAS_MAX_HEIGHT
|
|
max_width = video_width1
|
|
max_height = video_height1
|
|
duration1 = cut_video(input_file1, start_seconds1, end_seconds1, cut1, max_width, max_height, fps)
|
|
duration2 = cut_video(input_file2, start_seconds2, end_seconds2, cut2, max_width, max_height, fps)
|
|
max_duration = max(duration1, duration2)
|
|
if duration1 < max_duration:
|
|
loop_video_to_duration(cut1, cut1_looped, max_duration)
|
|
video1_file = cut1_looped
|
|
else:
|
|
video1_file = cut1
|
|
if duration2 < max_duration:
|
|
loop_video_to_duration(cut2, cut2_looped, max_duration)
|
|
video2_file = cut2_looped
|
|
else:
|
|
video2_file = cut2
|
|
overlay_new_width = int(max_width * 2 / 15)
|
|
overlay_new_height = int(max_width * 2 / 15)
|
|
overlay_x = (max_width * 2 - overlay_new_width) // 2
|
|
overlay_y = (max_height - overlay_new_height) // 2
|
|
process_videos_with_overlay(video1_file, video2_file, overlay_image, output_file,
|
|
overlay_new_width=overlay_new_width,
|
|
overlay_new_height=overlay_new_height,
|
|
overlay_x=overlay_x,
|
|
overlay_y=overlay_y,
|
|
)
|
|
if os.path.exists(cut1):
|
|
os.remove(cut1)
|
|
if os.path.exists(cut2):
|
|
os.remove(cut2)
|
|
if os.path.exists(cut1_looped):
|
|
os.remove(cut1_looped)
|
|
if os.path.exists(cut2_looped):
|
|
os.remove(cut2_looped)
|
|
def process_videos_with_overlay(input_file1, input_file2, overlay_image, output_file,
|
|
overlay_new_width=100, overlay_new_height=100,
|
|
overlay_x=0, overlay_y=0, fps=24):
|
|
|
|
command = [
|
|
ffmpeg_path_exe,
|
|
'-i', input_file1,
|
|
'-i', input_file2,
|
|
'-i', overlay_image,
|
|
'-filter_complex',
|
|
f"[0:v][1:v]hstack=inputs=2[base];"
|
|
f"[2:v]scale={overlay_new_width}:{overlay_new_height}[ovrl];"
|
|
f"[base][ovrl]overlay={overlay_x}:{overlay_y}",
|
|
'-c:v', 'libx264',
|
|
'-pix_fmt', 'yuv420p',
|
|
'-threads', '8',
|
|
'-y',
|
|
output_file
|
|
]
|
|
try:
|
|
subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"FFmpeg 命令执行失败: {e.stderr.decode()}")
|
|
raise
|
|
def apply_overlay_to_video(video_file, overlay_image_file, output_file, file_num=1):
|
|
|
|
video_width, video_height, size_mb = get_video_dimensions(video_file)
|
|
overlay_image = Image.open(overlay_image_file).convert("RGBA")
|
|
overlay_ratio = overlay_image.width / overlay_image.height
|
|
if file_num == 1:
|
|
if video_width < video_height:
|
|
overlay_new_width = video_width // 2
|
|
overlay_new_height = int(overlay_new_width / overlay_ratio)
|
|
else:
|
|
overlay_new_height = video_height // 2
|
|
overlay_new_width = int(overlay_new_height * overlay_ratio)
|
|
else:
|
|
overlay_new_width = int(video_width / 15)
|
|
overlay_new_height = int(video_width / 15)
|
|
overlay_resized = overlay_image.resize((overlay_new_width, overlay_new_height), resample_filter)
|
|
img_byte_arr = io.BytesIO()
|
|
overlay_resized.save(img_byte_arr, format='PNG')
|
|
img_byte_arr = img_byte_arr.getvalue()
|
|
overlay_x = (video_width - overlay_new_width) // 2
|
|
overlay_y = (video_height - overlay_new_height) // 2
|
|
command = [
|
|
ffmpeg_path_exe,
|
|
'-i', video_file,
|
|
'-i', 'pipe:0',
|
|
'-filter_complex',
|
|
f'[1:v]scale={overlay_new_width}:{overlay_new_height}[ovrl];[0:v][ovrl]overlay={overlay_x}:{overlay_y}',
|
|
'-c:v', 'libx264',
|
|
'-pix_fmt', 'yuv420p',
|
|
'-y',
|
|
output_file
|
|
]
|
|
process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
stdout, stderr = process.communicate(input=img_byte_arr)
|
|
if process.returncode != 0:
|
|
raise Exception(f"FFmpeg 命令执行失败: {stderr.decode()}")
|
|
def is_image(file_path):
|
|
|
|
if imghdr.what(file_path):
|
|
return True
|
|
mime_type, _ = mimetypes.guess_type(file_path)
|
|
if mime_type and mime_type.startswith('image'):
|
|
return mime_type
|
|
return False
|
|
def is_video(file_path):
|
|
|
|
mime_type, _ = mimetypes.guess_type(file_path)
|
|
if mime_type and mime_type.startswith('video'):
|
|
return mime_type
|
|
try:
|
|
command = [
|
|
ffmpeg_path_exe,
|
|
'-v', 'error',
|
|
'-select_streams', 'v:0',
|
|
'-show_entries', 'stream=codec_name',
|
|
'-of', 'csv=p=0',
|
|
'-i', file_path
|
|
]
|
|
result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
codec_name = result.stdout.strip()
|
|
if codec_name:
|
|
mime_types = {
|
|
'h264': 'video/mp4',
|
|
'hevc': 'video/mp4',
|
|
'vp8': 'video/webm',
|
|
'vp9': 'video/webm',
|
|
'av1': 'video/mp4',
|
|
'mpeg4': 'video/mp4',
|
|
'theora': 'video/ogg',
|
|
'wmv3': 'video/x-ms-wmv',
|
|
'vp6f': 'video/x-flv',
|
|
}
|
|
return mime_types.get(codec_name, 'video/unknown')
|
|
except FileNotFoundError:
|
|
print(f"{ffmpeg_path_exe} not found, unable to confirm if the file is a video.")
|
|
return False
|
|
CANVAS_MAX_WIDTH = 1000
|
|
CANVAS_MAX_HEIGHT = 1000
|
|
END_SECONDS = 3
|
|
def optimize_dimensions(w1, h1, w2, h2):
|
|
|
|
def adjust(value):
|
|
return value - (value % 4)
|
|
half_width = min(w2, CANVAS_MAX_WIDTH // 2)
|
|
height = int(half_width * h2 / w2)
|
|
return adjust(half_width * 2), adjust(height)
|
|
def process_file(file_path, side, target_width, target_height):
|
|
|
|
file_info = {
|
|
'file_type': '',
|
|
'mime_type': '',
|
|
'width': 0,
|
|
'height': 0,
|
|
'base_64': ''
|
|
}
|
|
file_mime_type, file_type = determine_file_type(file_path)
|
|
file_info['file_type'] = file_type
|
|
file_info['mime_type'] = file_mime_type
|
|
if file_type == 'image':
|
|
width, height, size_mb = get_image_dimensions(file_path)
|
|
if size_mb > 5:
|
|
compress_image(file_path, target_width=target_width)
|
|
width, height, _ = get_image_dimensions(file_path)
|
|
file_info['width'] = width
|
|
file_info['height'] = height
|
|
file_info['base_64'] = file_to_base64(file_path)
|
|
return file_path, file_info
|
|
elif file_type == 'video':
|
|
width, height, size_mb = get_video_dimensions(file_path)
|
|
if size_mb > 20:
|
|
cut_and_compress_video(file_path, target_width=target_width, duration=5)
|
|
width, height, _ = get_video_dimensions(file_path)
|
|
file_info['width'] = width
|
|
file_info['height'] = height
|
|
file_info['base_64'] = file_to_base64(file_path)
|
|
return file_path, file_info
|
|
else:
|
|
raise ValueError(f"Unsupported file type for {side} side.")
|
|
def do_zhutu(file_left='', file_right='', overlay_path=''):
|
|
|
|
if not file_left and not file_right:
|
|
return {'error': '请上传两张图片或两张视频', 'code': 1}
|
|
out_put_name = generate_md5_uid_timestamp(overlay_path)
|
|
out_put_png = out_put_name + '.png'
|
|
out_put_mp4 = out_put_name + '.mp4'
|
|
out_put_mp4_temp = out_put_name + 'temp.mp4'
|
|
out_put_name_mp4_temp = folder_paths.get_input_directory() + '/' + out_put_mp4_temp
|
|
out_put_name_mp4 = folder_paths.get_input_directory() + '/' + out_put_mp4
|
|
out_put_name_png = folder_paths.get_input_directory() + '/' + out_put_png
|
|
base_size = {}
|
|
with ThreadPoolExecutor() as executor:
|
|
futures = {}
|
|
if file_left:
|
|
futures['left'] = executor.submit(process_file, file_left, 'left', CANVAS_MAX_WIDTH // 2, CANVAS_MAX_HEIGHT)
|
|
if file_right:
|
|
futures['right'] = executor.submit(process_file, file_right, 'right', CANVAS_MAX_WIDTH // 2, CANVAS_MAX_HEIGHT)
|
|
for side, future in futures.items():
|
|
try:
|
|
file_path, file_info = future.result()
|
|
base_size[side] = file_info
|
|
except Exception as e:
|
|
return {'error': str(e), 'code': 1}
|
|
if 'left' in base_size and 'right' in base_size:
|
|
canvas_width, canvas_height = optimize_dimensions(
|
|
base_size['left']['width'], base_size['left']['height'],
|
|
base_size['right']['width'], base_size['right']['height']
|
|
)
|
|
else:
|
|
canvas_width, canvas_height = CANVAS_MAX_WIDTH, CANVAS_MAX_HEIGHT
|
|
try:
|
|
if 'left' in base_size and 'right' in base_size:
|
|
if base_size['left']['file_type'] == 'image' and base_size['right']['file_type'] == 'image':
|
|
process_image(
|
|
input_image_path_left=file_left,
|
|
input_image_path_right=file_right,
|
|
output_image_path=out_put_name_png,
|
|
canvas_size=(canvas_width, canvas_height),
|
|
overlay_image_path=overlay_path
|
|
)
|
|
result_type = 'image'
|
|
elif base_size['left']['file_type'] == 'video' and base_size['right']['file_type'] == 'video':
|
|
process_and_merge_videos(
|
|
file_left, file_right, 0, END_SECONDS, 0, END_SECONDS,
|
|
output_file=out_put_name_mp4,
|
|
overlay_image=overlay_path
|
|
)
|
|
result_type = 'video'
|
|
else:
|
|
if base_size['left']['file_type'] == 'video':
|
|
process_and_merge_image_video(
|
|
file_right, file_left, out_put_name_mp4,
|
|
overlay_path,
|
|
side='right', target_width=canvas_width // 2, target_height=canvas_height,
|
|
start_seconds=0, end_seconds=END_SECONDS, fps=24
|
|
)
|
|
else:
|
|
process_and_merge_image_video(
|
|
file_left, file_right, out_put_name_mp4,
|
|
overlay_path,
|
|
side='left', target_width=canvas_width // 2, target_height=canvas_height,
|
|
start_seconds=0, end_seconds=END_SECONDS, fps=24
|
|
)
|
|
result_type = 'video'
|
|
else:
|
|
if 'left' in base_size:
|
|
single_file = file_left
|
|
side = 'left'
|
|
file_info = base_size['left']
|
|
else:
|
|
single_file = file_right
|
|
side = 'right'
|
|
file_info = base_size['right']
|
|
if file_info['file_type'] == 'image':
|
|
process_image(
|
|
**{f"input_image_path_{side}": single_file},
|
|
output_image_path=out_put_name_png,
|
|
canvas_size=(canvas_width, canvas_height),
|
|
overlay_image_path=overlay_path
|
|
)
|
|
result_type = 'image'
|
|
else:
|
|
process_and_merge_image_video(
|
|
None, single_file, out_put_name_mp4,
|
|
overlay_path,
|
|
side='right' if side == 'left' else 'left',
|
|
target_width=canvas_width // 2, target_height=canvas_height,
|
|
start_seconds=0, end_seconds=END_SECONDS, fps=24
|
|
)
|
|
result_type = 'video'
|
|
if result_type == 'image':
|
|
result_data = file_to_base64(out_put_name_png)
|
|
filename = out_put_png
|
|
else:
|
|
result_data = file_to_base64(out_put_name_mp4)
|
|
filename = out_put_mp4
|
|
return {
|
|
'data': result_data,
|
|
'filename': filename,
|
|
'code': 0,
|
|
'type': result_type,
|
|
'size': {'width': canvas_width, 'height': canvas_height},
|
|
'base_size': base_size
|
|
}
|
|
finally:
|
|
if os.path.exists(out_put_name_mp4_temp):
|
|
os.remove(out_put_name_mp4_temp)
|
|
def get_video_duration(video_path):
|
|
result = subprocess.run(
|
|
[ffprobe_exe_path, '-v', 'error', '-show_entries', 'format=duration',
|
|
'-of', 'default=noprint_wrappers=1:nokey=1', video_path],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True
|
|
)
|
|
return float(result.stdout.strip())
|
|
def extract_frames(video_path, num_frames=3):
|
|
duration = get_video_duration(video_path)
|
|
frame_times = []
|
|
if duration < 1.0:
|
|
frame_times = [0]
|
|
else:
|
|
frame_times = [(duration / (num_frames + 1)) * (i + 1) for i in range(num_frames)]
|
|
frame_contents = []
|
|
for time in frame_times:
|
|
result = subprocess.run(
|
|
[ffmpeg_path_exe, '-ss', str(time), '-i', video_path, '-vframes', '1',
|
|
'-f', 'image2pipe', '-vcodec', 'png', '-'],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE
|
|
)
|
|
frame_contents.append(result.stdout)
|
|
return frame_contents
|