| 1 | # -*- coding: utf-8 -*- |
| 2 | """ |
| 3 | 阶段6: 后期制作智能体 |
| 4 | 拼接用户在阶段5选定的视频片段 → 最终成片 |
| 5 | """ |
| 6 | |
| 7 | import os |
| 8 | import re |
| 9 | import asyncio |
| 10 | import logging |
| 11 | import subprocess |
| 12 | from typing import Any, Optional, Dict |
| 13 | |
| 14 | from .base_agent import AgentInterface |
| 15 | |
| 16 | logger = logging.getLogger(__name__) |
| 17 | |
| 18 | |
| 19 | class VideoEditorAgent(AgentInterface): |
| 20 | """后期制作：拼接用户选择的视频片段 → 最终成片""" |
| 21 | |
| 22 | def __init__(self): |
| 23 | super().__init__(name="VideoEditor") |
| 24 | |
| 25 | async def process(self, input_data: Any, intervention: Optional[Dict] = None) -> Dict: |
| 26 | input_data = self._merge_session_params(input_data) |
| 27 | sid = input_data["session_id"] |
| 28 | regenerate_episodes = set() |
| 29 | if isinstance(intervention, dict) and isinstance(intervention.get("regenerate_episodes"), list): |
| 30 | regenerate_episodes = { |
| 31 | int(ep) |
| 32 | for ep in intervention.get("regenerate_episodes", []) |
| 33 | if str(ep).isdigit() |
| 34 | } |
| 35 | |
| 36 | # 从编排器注入的统一 session 快照读取阶段5数据。 |
| 37 | artifacts = self._session_artifacts(input_data) |
| 38 | video_art = artifacts.get("video_generation", {}) |
| 39 | clips_list = video_art.get("clips", []) |
| 40 | post_art = artifacts.get("post_production", {}) |
| 41 | |
| 42 | # 获取剧集标题映射 (从 Storyboard) |
| 43 | storyboard_art = artifacts.get("storyboard", {}) |
| 44 | episodes_info = storyboard_art.get("episodes", []) |
| 45 | ep_title_map = {int(ep.get("episode_number", 0)): ep.get("act_title", f"第 {ep.get('episode_number')} 集") |
| 46 | for ep in episodes_info if ep.get("episode_number")} |
| 47 | |
| 48 | if not clips_list: |
| 49 | # 兼容旧逻辑：如果 artifacts 里没有，尝试从 input_data 获取 |
| 50 | selected_clips: dict = input_data.get("selected_clips", {}) |
| 51 | if not selected_clips: |
| 52 | raise Exception("未找到选定的视频片段数据，请先完成阶段5") |
| 53 | |
| 54 | self._report_progress("后期制作", "准备视频片段...", 5) |
| 55 | |
| 56 | def run(): |
| 57 | video_dir = os.path.join('code/result/video', str(sid)) |
| 58 | os.makedirs(video_dir, exist_ok=True) |
| 59 | output_dir = os.path.join(video_dir, 'output') |
| 60 | os.makedirs(output_dir, exist_ok=True) |
| 61 | |
| 62 | # 按剧集分组片段 |
| 63 | episodes_map = {} # { episode_index: [clip_paths] } |
| 64 | |
| 65 | if clips_list: |
| 66 | for clip in clips_list: |
| 67 | path = clip.get("selected") |
| 68 | if not path or not os.path.exists(path): |
| 69 | logger.warning(f"[{sid}] Clip missing: {clip.get('id')} → {path}") |
| 70 | continue |
| 71 | |
| 72 | # 优先使用片段数据中的 episode 字段 |
| 73 | ep_idx = clip.get("episode") |
| 74 | |
| 75 | # 如果没有 episode 字段，则尝试通过 ID 解析 |
| 76 | if ep_idx is None: |
| 77 | match = re.search(r'(?:seg_|shot_)?(\d{1,3})_\d{1,3}', clip.get("id", "")) |
| 78 | if match: |
| 79 | ep_idx = int(match.group(1)) |
| 80 | else: |
| 81 | # 最后的归底方案：从 name 提取或默认为 1 |
| 82 | name_match = re.search(r'第(\d+)集', clip.get("name", "")) |
| 83 | if not name_match: |
| 84 | name_match = re.search(r'(\d+)', clip.get("name", "")) |
| 85 | ep_idx = int(name_match.group(1)) if name_match else 1 |
| 86 | |
| 87 | if regenerate_episodes and int(ep_idx) not in regenerate_episodes: |
| 88 | continue |
| 89 | |
| 90 | episodes_map.setdefault(int(ep_idx), []).append(path) |
| 91 | else: |
| 92 | # 兼容旧逻辑 |
| 93 | def sort_key(k: str) -> tuple: |
| 94 | return tuple(int(n) for n in re.findall(r'\d+', k)) or (999,) |
| 95 | selected_clips = input_data.get("selected_clips", {}) |
| 96 | for shot_id in sorted(selected_clips.keys(), key=sort_key): |
| 97 | path = selected_clips[shot_id] |
| 98 | if os.path.exists(path): |
| 99 | # 旧逻辑默认全部归为第1集 |
| 100 | if not regenerate_episodes or 1 in regenerate_episodes: |
| 101 | episodes_map.setdefault(1, []).append(path) |
| 102 | else: |
| 103 | logger.warning(f"[{sid}] Clip missing: {shot_id} → {path}") |
| 104 | |
| 105 | if not episodes_map: |
| 106 | raise Exception("没有可用于拼接的视频文件") |
| 107 | |
| 108 | final_videos = [] |
| 109 | sorted_episodes = sorted(episodes_map.keys()) |
| 110 | total_eps = len(sorted_episodes) |
| 111 | |
| 112 | ffmpeg_exe = 'ffmpeg' |
| 113 | |
| 114 | for i, ep_idx in enumerate(sorted_episodes): |
| 115 | self._report_progress("后期制作", f"正在拼接第 {ep_idx} 集 ({i+1}/{total_eps})...", int(20 + (i/total_eps)*70)) |
| 116 | |
| 117 | clip_paths = episodes_map[ep_idx] |
| 118 | list_file = os.path.join(video_dir, f'concat_list_ep{ep_idx}.txt') |
| 119 | output = os.path.join(output_dir, f'{sid}_ep{ep_idx}.mp4') |
| 120 | |
| 121 | with open(list_file, 'w', encoding='utf-8') as f: |
| 122 | for p in clip_paths: |
| 123 | abs_p = os.path.abspath(p).replace('\\', '/') |
| 124 | f.write(f"file '{abs_p}'\n") |
| 125 | |
| 126 | cmd = [ |
| 127 | ffmpeg_exe, '-y', '-f', 'concat', '-safe', '0', |
| 128 | '-i', list_file, |
| 129 | '-c:v', 'libx264', '-preset', 'fast', '-crf', '22', |
| 130 | '-c:a', 'aac', '-pix_fmt', 'yuv420p', |
| 131 | '-movflags', '+faststart', output |
| 132 | ] |
| 133 | |
| 134 | logger.info(f"[{sid}] Running ffmpeg for Ep {ep_idx}: {cmd}") |
| 135 | try: |
| 136 | result = subprocess.run(cmd, capture_output=True, text=True, check=True) |
| 137 | except subprocess.CalledProcessError as e: |
| 138 | logger.error(f"FFmpeg failed with exit code {e.returncode}") |
| 139 | logger.error(f"FFmpeg stderr: {e.stderr}") |
| 140 | raise Exception(f"视频拼接失败: {e.stderr}") |
| 141 | |
| 142 | ep_title = ep_title_map.get(ep_idx, f"第 {ep_idx} 集") |
| 143 | final_videos.append({ |
| 144 | "episode": ep_idx, |
| 145 | "path": output, |
| 146 | "name": ep_title |
| 147 | }) |
| 148 | |
| 149 | return final_videos |
| 150 | |
| 151 | loop = asyncio.get_running_loop() |
| 152 | final_results = await loop.run_in_executor(None, run) |
| 153 | |
| 154 | if regenerate_episodes and isinstance(post_art, dict): |
| 155 | existing_videos = post_art.get("final_videos", []) |
| 156 | if isinstance(existing_videos, list): |
| 157 | regenerated_eps = {item.get("episode") for item in final_results if isinstance(item, dict)} |
| 158 | preserved = [ |
| 159 | item for item in existing_videos |
| 160 | if isinstance(item, dict) and item.get("episode") not in regenerated_eps |
| 161 | ] |
| 162 | final_results = sorted( |
| 163 | preserved + final_results, |
| 164 | key=lambda item: item.get("episode", 999) if isinstance(item, dict) else 999, |
| 165 | ) |
| 166 | |
| 167 | self._report_progress("后期制作", "成片完成", 100) |
| 168 | |
| 169 | return { |
| 170 | "payload": { |
| 171 | "session_id": sid, |
| 172 | "final_videos": final_results, |
| 173 | "final_video": final_results[0]["path"] if final_results else "", |
| 174 | }, |
| 175 | "stage_completed": True, |
| 176 | } |
| 177 |