Initial clean project import
This commit is contained in:
cat-shark
@@ -0,0 +1,530 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
字幕封面生成器 - 一键生成视频字幕和封面
|
||||
Subtitle & Cover Generator for AIGC Panel
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import argparse
|
||||
from typing import Dict, List, Any, Tuple
|
||||
from pathlib import Path
|
||||
|
||||
# 导入必要的库
|
||||
try:
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw, ImageFont, ImageFilter, ImageEnhance, ImageOps
|
||||
from rembg import remove
|
||||
except ImportError as e:
|
||||
print(f"Error: Missing required library - {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
class SubtitleCoverGenerator:
|
||||
"""字幕封面生成器主类"""
|
||||
|
||||
def __init__(self, config: Dict[str, Any]):
|
||||
"""
|
||||
初始化生成器
|
||||
|
||||
Args:
|
||||
config: 配置字典，包含视频路径、文案、样式等
|
||||
"""
|
||||
self.config = config
|
||||
self.video_path = config.get('videoPath')
|
||||
self.script_text = config.get('scriptText', '')
|
||||
self.subtitle_style = config.get('subtitleStyle', 'default')
|
||||
self.cover_style = config.get('coverStyle', 'default')
|
||||
self.output_dir = config.get('outputDir', './output')
|
||||
|
||||
# 新增配置参数
|
||||
self.is_portrait = config.get('isPortrait', False) # 是否为竖屏视频
|
||||
self.cover_text = config.get('coverText', '') # 封面文字
|
||||
self.cover_text_position = config.get('coverTextPosition', 'bottom') # 文字位置: top, center, bottom
|
||||
self.cover_text_color = config.get('coverTextColor', '#FFFFFF') # 文字颜色
|
||||
self.cover_text_size = config.get('coverTextSize', 50) # 文字大小
|
||||
self.outline_color = config.get('outlineColor', '#FFD700') # 描边颜色
|
||||
self.outline_width = config.get('outlineWidth', 3) # 描边宽度
|
||||
self.blur_radius = config.get('blurRadius', 15) # 背景模糊度
|
||||
self.zoom_level = config.get('zoomLevel', 1.0) # 封面放大倍数
|
||||
|
||||
# 确保输出目录存在
|
||||
os.makedirs(self.output_dir, exist_ok=True)
|
||||
|
||||
def get_video_info(self) -> Dict[str, Any]:
|
||||
"""
|
||||
获取视频信息
|
||||
|
||||
Returns:
|
||||
包含视频宽度、高度、时长等信息的字典
|
||||
"""
|
||||
try:
|
||||
cap = cv2.VideoCapture(self.video_path)
|
||||
if not cap.isOpened():
|
||||
raise Exception("Could not open video file")
|
||||
|
||||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
duration = frame_count / fps if fps > 0 else 0
|
||||
|
||||
cap.release()
|
||||
|
||||
# 检测是否为竖屏视频
|
||||
is_portrait = height > width
|
||||
|
||||
return {
|
||||
'width': width,
|
||||
'height': height,
|
||||
'fps': fps,
|
||||
'duration': duration,
|
||||
'is_portrait': is_portrait
|
||||
}
|
||||
except Exception as e:
|
||||
raise Exception(f"Error getting video info: {str(e)}")
|
||||
|
||||
def extract_frame(self, timestamp: float = 1.0) -> np.ndarray:
|
||||
"""
|
||||
从视频中提取关键帧
|
||||
|
||||
Args:
|
||||
timestamp: 提取帧的时间戳（秒）
|
||||
|
||||
Returns:
|
||||
提取的帧（numpy数组）
|
||||
"""
|
||||
try:
|
||||
cap = cv2.VideoCapture(self.video_path)
|
||||
|
||||
# 设置到指定时间戳
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
frame_number = int(timestamp * fps)
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
|
||||
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
|
||||
if not ret:
|
||||
raise Exception("Failed to extract frame from video")
|
||||
|
||||
return frame
|
||||
except Exception as e:
|
||||
raise Exception(f"Error extracting frame: {str(e)}")
|
||||
|
||||
def remove_background(self, image: np.ndarray) -> Image.Image:
|
||||
"""
|
||||
使用rembg进行人物抠图
|
||||
|
||||
Args:
|
||||
image: 输入图像（numpy数组）
|
||||
|
||||
Returns:
|
||||
抠图后的PIL Image对象（带透明通道）
|
||||
"""
|
||||
try:
|
||||
# 转换为PIL Image
|
||||
pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||||
|
||||
# 使用rembg移除背景
|
||||
output = remove(pil_image)
|
||||
|
||||
return output
|
||||
except Exception as e:
|
||||
raise Exception(f"Error removing background: {str(e)}")
|
||||
|
||||
def generate_cover(self) -> str:
|
||||
"""
|
||||
生成视频封面
|
||||
|
||||
Returns:
|
||||
生成的封面图片路径
|
||||
"""
|
||||
try:
|
||||
# 获取视频信息
|
||||
video_info = self.get_video_info()
|
||||
self.is_portrait = video_info['is_portrait']
|
||||
|
||||
# 提取关键帧
|
||||
frame = self.extract_frame(1.0)
|
||||
|
||||
if self.cover_style == 'default':
|
||||
# 默认封面：直接使用提取的帧
|
||||
cover_path = os.path.join(self.output_dir, 'cover.png')
|
||||
cv2.imwrite(cover_path, frame)
|
||||
|
||||
elif self.cover_style == 'blur-bg':
|
||||
# 背景虚化效果
|
||||
cover_path = self._generate_blur_cover(frame)
|
||||
|
||||
elif self.cover_style == 'outline':
|
||||
# 人物描边效果
|
||||
cover_path = self._generate_outline_cover(frame)
|
||||
|
||||
elif self.cover_style == 'gradient':
|
||||
# 渐变背景效果
|
||||
cover_path = self._generate_gradient_cover(frame)
|
||||
|
||||
elif self.cover_style == 'split':
|
||||
# 斜切设计效果
|
||||
cover_path = self._generate_split_cover(frame)
|
||||
else:
|
||||
cover_path = os.path.join(self.output_dir, 'cover.png')
|
||||
cv2.imwrite(cover_path, frame)
|
||||
|
||||
# 如果需要放大封面
|
||||
if self.zoom_level != 1.0:
|
||||
cover_path = self._zoom_cover(cover_path)
|
||||
|
||||
# 如果有文字需要添加
|
||||
if self.cover_text.strip():
|
||||
cover_path = self._add_text_to_cover(cover_path)
|
||||
|
||||
return cover_path
|
||||
except Exception as e:
|
||||
raise Exception(f"Error generating cover: {str(e)}")
|
||||
|
||||
def _zoom_cover(self, cover_path: str) -> str:
|
||||
"""放大封面"""
|
||||
try:
|
||||
image = Image.open(cover_path)
|
||||
|
||||
# 计算新的尺寸
|
||||
new_width = int(image.width * self.zoom_level)
|
||||
new_height = int(image.height * self.zoom_level)
|
||||
|
||||
# 放大图像
|
||||
zoomed = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
# 如果是竖屏视频，需要裁剪中心区域
|
||||
if self.is_portrait:
|
||||
# 裁剪为原始尺寸的正方形区域
|
||||
left = (zoomed.width - image.width) // 2
|
||||
top = (zoomed.height - image.height) // 2
|
||||
right = left + image.width
|
||||
bottom = top + image.height
|
||||
zoomed = zoomed.crop((left, top, right, bottom))
|
||||
else:
|
||||
# 横屏视频裁剪中心区域
|
||||
left = (zoomed.width - image.width) // 2
|
||||
top = (zoomed.height - image.height) // 2
|
||||
right = left + image.width
|
||||
bottom = top + image.height
|
||||
zoomed = zoomed.crop((left, top, right, bottom))
|
||||
|
||||
# 保存放大后的封面
|
||||
zoom_path = os.path.join(self.output_dir, 'cover_zoomed.png')
|
||||
zoomed.save(zoom_path, quality=95)
|
||||
|
||||
return zoom_path
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not zoom cover: {e}", file=sys.stderr)
|
||||
return cover_path
|
||||
|
||||
def _add_text_to_cover(self, cover_path: str) -> str:
|
||||
"""添加文字到封面"""
|
||||
try:
|
||||
image = Image.open(cover_path)
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
# 尝试加载字体，如果失败则使用默认字体
|
||||
try:
|
||||
font = ImageFont.truetype("arial.ttf", self.cover_text_size)
|
||||
except:
|
||||
try:
|
||||
font = ImageFont.truetype("DejaVuSans-Bold.ttf", self.cover_text_size)
|
||||
except:
|
||||
font = ImageFont.load_default()
|
||||
|
||||
# 计算文字位置
|
||||
bbox = draw.textbbox((0, 0), self.cover_text, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
text_height = bbox[3] - bbox[1]
|
||||
|
||||
if self.cover_text_position == 'top':
|
||||
x = (image.width - text_width) // 2
|
||||
y = 50
|
||||
elif self.cover_text_position == 'center':
|
||||
x = (image.width - text_width) // 2
|
||||
y = (image.height - text_height) // 2
|
||||
else: # bottom
|
||||
x = (image.width - text_width) // 2
|
||||
y = image.height - text_height - 50
|
||||
|
||||
# 添加描边效果
|
||||
if self.outline_width > 0:
|
||||
# 创建描边
|
||||
outline_color = self._hex_to_rgb(self.outline_color)
|
||||
for offset_x in range(-self.outline_width, self.outline_width + 1):
|
||||
for offset_y in range(-self.outline_width, self.outline_width + 1):
|
||||
if offset_x == 0 and offset_y == 0:
|
||||
continue
|
||||
draw.text((x + offset_x, y + offset_y), self.cover_text, fill=outline_color, font=font)
|
||||
|
||||
# 添加文字
|
||||
text_color = self._hex_to_rgb(self.cover_text_color)
|
||||
draw.text((x, y), self.cover_text, fill=text_color, font=font)
|
||||
|
||||
# 保存带文字的封面
|
||||
text_path = os.path.join(self.output_dir, 'cover_with_text.png')
|
||||
image.save(text_path, quality=95)
|
||||
|
||||
return text_path
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not add text to cover: {e}", file=sys.stderr)
|
||||
return cover_path
|
||||
|
||||
def _hex_to_rgb(self, hex_color: str) -> Tuple[int, int, int]:
|
||||
"""将十六进制颜色转换为RGB元组"""
|
||||
hex_color = hex_color.lstrip('#')
|
||||
return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
|
||||
|
||||
def _generate_blur_cover(self, frame: np.ndarray) -> str:
|
||||
"""生成背景虚化封面"""
|
||||
# 转换为PIL
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
|
||||
# 抠出人物
|
||||
person = self.remove_background(frame)
|
||||
|
||||
# 背景虚化（使用配置的模糊度）
|
||||
background = image.filter(ImageFilter.GaussianBlur(radius=self.blur_radius))
|
||||
|
||||
# 合成
|
||||
background.paste(person, (0, 0), person)
|
||||
|
||||
# 保存
|
||||
cover_path = os.path.join(self.output_dir, 'cover_blur.png')
|
||||
background.save(cover_path, quality=95)
|
||||
|
||||
return cover_path
|
||||
|
||||
def _generate_outline_cover(self, frame: np.ndarray) -> str:
|
||||
"""生成人物描边封面"""
|
||||
# 转换为PIL
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
|
||||
# 抠出人物
|
||||
person = self.remove_background(frame)
|
||||
|
||||
# 创建描边效果
|
||||
# 将person转为numpy进行描边处理
|
||||
person_np = np.array(person)
|
||||
alpha = person_np[:, :, 3]
|
||||
|
||||
# 使用形态学操作创建描边
|
||||
kernel = np.ones((self.outline_width * 2 + 1, self.outline_width * 2 + 1), np.uint8)
|
||||
dilated = cv2.dilate(alpha, kernel, iterations=1)
|
||||
outline = dilated - alpha
|
||||
|
||||
# 创建描边图层
|
||||
outline_layer = Image.new('RGBA', person.size, (0, 0, 0, 0))
|
||||
outline_draw = ImageDraw.Draw(outline_layer)
|
||||
|
||||
# 获取描边颜色
|
||||
outline_color = self._hex_to_rgb(self.outline_color)
|
||||
|
||||
# 应用描边
|
||||
for y in range(outline.shape[0]):
|
||||
for x in range(outline.shape[1]):
|
||||
if outline[y, x] > 0:
|
||||
outline_layer.putpixel((x, y), outline_color + (255,))
|
||||
|
||||
# 合成
|
||||
result = Image.new('RGBA', image.size)
|
||||
result.paste(image, (0, 0))
|
||||
result.paste(outline_layer, (0, 0), outline_layer)
|
||||
result.paste(person, (0, 0), person)
|
||||
|
||||
# 保存
|
||||
cover_path = os.path.join(self.output_dir, 'cover_outline.png')
|
||||
result.save(cover_path, quality=95)
|
||||
|
||||
return cover_path
|
||||
|
||||
def _generate_gradient_cover(self, frame: np.ndarray) -> str:
|
||||
"""生成渐变背景封面"""
|
||||
# 转换为PIL
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
|
||||
# 抠出人物
|
||||
person = self.remove_background(frame)
|
||||
|
||||
# 创建渐变背景
|
||||
width, height = image.size
|
||||
gradient = Image.new('RGB', (width, height), color=(0, 0, 0))
|
||||
draw = ImageDraw.Draw(gradient)
|
||||
|
||||
# 绘制渐变
|
||||
for y in range(height):
|
||||
r = int(103 + (118 - 103) * y / height)
|
||||
g = int(126 + (74 - 126) * y / height)
|
||||
b = int(234 + (162 - 234) * y / height)
|
||||
draw.line([(0, y), (width, y)], fill=(r, g, b))
|
||||
|
||||
# 合成
|
||||
gradient.paste(person, (0, 0), person)
|
||||
|
||||
# 保存
|
||||
cover_path = os.path.join(self.output_dir, 'cover_gradient.png')
|
||||
gradient.save(cover_path, quality=95)
|
||||
|
||||
return cover_path
|
||||
|
||||
def _generate_split_cover(self, frame: np.ndarray) -> str:
|
||||
"""生成斜切设计封面"""
|
||||
# 转换为PIL
|
||||
image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
width, height = image.size
|
||||
|
||||
# 创建斜切背景
|
||||
background = Image.new('RGB', (width, height), color=(255, 255, 255))
|
||||
draw = ImageDraw.Draw(background)
|
||||
|
||||
# 绘制斜切色块
|
||||
draw.polygon([(0, 0), (width * 0.6, 0), (width * 0.4, height), (0, height)],
|
||||
fill=(103, 126, 234))
|
||||
draw.polygon([(width * 0.6, 0), (width, 0), (width, height), (width * 0.4, height)],
|
||||
fill=(118, 74, 162))
|
||||
|
||||
# 应用半透明图片
|
||||
image_enhanced = ImageEnhance.Brightness(image).enhance(0.8)
|
||||
background.paste(image_enhanced, (0, 0))
|
||||
|
||||
# 保存
|
||||
cover_path = os.path.join(self.output_dir, 'cover_split.png')
|
||||
background.save(cover_path, quality=95)
|
||||
|
||||
return cover_path
|
||||
|
||||
def generate_subtitles(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
生成字幕列表
|
||||
|
||||
Returns:
|
||||
字幕列表，每个字幕包含文本、开始时间、结束时间
|
||||
"""
|
||||
try:
|
||||
# 将文案按行分割
|
||||
lines = [line.strip() for line in self.script_text.split('\n') if line.strip()]
|
||||
|
||||
# 生成字幕列表（简单示例，每句3秒）
|
||||
subtitles = []
|
||||
for i, line in enumerate(lines):
|
||||
subtitle = {
|
||||
'index': i + 1,
|
||||
'text': line,
|
||||
'start_time': i * 3.0,
|
||||
'end_time': (i + 1) * 3.0,
|
||||
'style': self.subtitle_style
|
||||
}
|
||||
subtitles.append(subtitle)
|
||||
|
||||
return subtitles
|
||||
except Exception as e:
|
||||
raise Exception(f"Error generating subtitles: {str(e)}")
|
||||
|
||||
def save_srt(self, subtitles: List[Dict[str, Any]]) -> str:
|
||||
"""
|
||||
保存SRT字幕文件
|
||||
|
||||
Args:
|
||||
subtitles: 字幕列表
|
||||
|
||||
Returns:
|
||||
SRT文件路径
|
||||
"""
|
||||
try:
|
||||
srt_path = os.path.join(self.output_dir, 'subtitles.srt')
|
||||
|
||||
with open(srt_path, 'w', encoding='utf-8') as f:
|
||||
for sub in subtitles:
|
||||
# SRT格式
|
||||
f.write(f"{sub['index']}\n")
|
||||
f.write(f"{self._format_time(sub['start_time'])} --> {self._format_time(sub['end_time'])}\n")
|
||||
f.write(f"{sub['text']}\n\n")
|
||||
|
||||
return srt_path
|
||||
except Exception as e:
|
||||
raise Exception(f"Error saving SRT file: {str(e)}")
|
||||
|
||||
def _format_time(self, seconds: float) -> str:
|
||||
"""格式化时间为SRT格式"""
|
||||
hours = int(seconds // 3600)
|
||||
minutes = int((seconds % 3600) // 60)
|
||||
secs = int(seconds % 60)
|
||||
millis = int((seconds % 1) * 1000)
|
||||
|
||||
return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
|
||||
|
||||
def generate_all(self) -> Dict[str, Any]:
|
||||
"""
|
||||
一键生成所有内容
|
||||
|
||||
Returns:
|
||||
包含生成结果的字典
|
||||
"""
|
||||
try:
|
||||
# 生成封面
|
||||
print("Generating cover...", file=sys.stderr)
|
||||
cover_path = self.generate_cover()
|
||||
|
||||
# 生成字幕
|
||||
print("Generating subtitles...", file=sys.stderr)
|
||||
subtitles = self.generate_subtitles()
|
||||
srt_path = self.save_srt(subtitles)
|
||||
|
||||
result = {
|
||||
'success': True,
|
||||
'coverPath': cover_path,
|
||||
'srtPath': srt_path,
|
||||
'subtitles': subtitles,
|
||||
'message': 'Generated successfully'
|
||||
}
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': str(e),
|
||||
'message': f'Generation failed: {str(e)}'
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
parser = argparse.ArgumentParser(description='Subtitle & Cover Generator')
|
||||
parser.add_argument('config_file', help='Configuration JSON file path')
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
# 读取配置文件
|
||||
with open(args.config_file, 'r', encoding='utf-8') as f:
|
||||
config = json.load(f)
|
||||
|
||||
# 创建生成器
|
||||
generator = SubtitleCoverGenerator(config)
|
||||
|
||||
# 执行生成
|
||||
result = generator.generate_all()
|
||||
|
||||
# 输出结果
|
||||
print(json.dumps(result, ensure_ascii=False))
|
||||
|
||||
sys.exit(0 if result['success'] else 1)
|
||||
|
||||
except Exception as e:
|
||||
error_result = {
|
||||
'success': False,
|
||||
'error': str(e),
|
||||
'message': f'Error: {str(e)}'
|
||||
}
|
||||
print(json.dumps(error_result, ensure_ascii=False))
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user