306 lines
11 KiB
Python
306 lines
11 KiB
Python
"""
|
|
Video assembler module for concatenating processed segments.
|
|
Handles merging processed segments and adding audio from original video.
|
|
"""
|
|
|
|
import os
|
|
import subprocess
|
|
import logging
|
|
from typing import List, Optional
|
|
from utils.file_utils import get_segments_directories, file_exists
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class VideoAssembler:
|
|
"""Handles final video assembly from processed segments."""
|
|
|
|
def __init__(self, preserve_audio: bool = True, use_nvenc: bool = False,
|
|
output_mode: str = "green_screen"):
|
|
"""
|
|
Initialize video assembler.
|
|
|
|
Args:
|
|
preserve_audio: Whether to preserve audio from original video
|
|
use_nvenc: Whether to use hardware encoding for final output
|
|
output_mode: Output mode - "green_screen" or "alpha_channel"
|
|
"""
|
|
self.preserve_audio = preserve_audio
|
|
self.use_nvenc = use_nvenc
|
|
self.output_mode = output_mode
|
|
|
|
def create_concat_file(self, segments_dir: str, output_filename: str = "concat_list.txt") -> Optional[str]:
|
|
"""
|
|
Create a concatenation file for FFmpeg.
|
|
|
|
Args:
|
|
segments_dir: Directory containing processed segments
|
|
output_filename: Name for the concat file
|
|
|
|
Returns:
|
|
Path to concat file or None if no valid segments found
|
|
"""
|
|
concat_path = os.path.join(segments_dir, output_filename)
|
|
valid_segments = 0
|
|
|
|
try:
|
|
segments = get_segments_directories(segments_dir)
|
|
|
|
with open(concat_path, 'w') as f:
|
|
for i, segment in enumerate(segments):
|
|
segment_dir = os.path.join(segments_dir, segment)
|
|
if self.output_mode == "alpha_channel":
|
|
output_video = os.path.join(segment_dir, f"output_{i}.mov")
|
|
else:
|
|
output_video = os.path.join(segment_dir, f"output_{i}.mp4")
|
|
|
|
if file_exists(output_video):
|
|
# Use relative path for FFmpeg
|
|
relative_path = os.path.relpath(output_video, segments_dir)
|
|
f.write(f"file '{relative_path}'\n")
|
|
valid_segments += 1
|
|
else:
|
|
logger.warning(f"Output video not found for segment {i}: {output_video}")
|
|
|
|
if valid_segments == 0:
|
|
logger.error("No valid output segments found for concatenation")
|
|
os.remove(concat_path)
|
|
return None
|
|
|
|
logger.info(f"Created concatenation file with {valid_segments} segments: {concat_path}")
|
|
return concat_path
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error creating concatenation file: {e}")
|
|
return None
|
|
|
|
def concatenate_segments(self, segments_dir: str, output_path: str,
|
|
bitrate: str = "50M") -> bool:
|
|
"""
|
|
Concatenate video segments using FFmpeg.
|
|
|
|
Args:
|
|
segments_dir: Directory containing processed segments
|
|
output_path: Path for final concatenated video
|
|
bitrate: Output video bitrate
|
|
|
|
Returns:
|
|
True if successful
|
|
"""
|
|
# Create concatenation file
|
|
concat_file = self.create_concat_file(segments_dir)
|
|
if not concat_file:
|
|
return False
|
|
|
|
try:
|
|
# Build FFmpeg command
|
|
if self.output_mode == "alpha_channel":
|
|
# For alpha channel, we need to maintain the ProRes codec
|
|
cmd = [
|
|
'ffmpeg',
|
|
'-y', # Overwrite output
|
|
'-f', 'concat',
|
|
'-safe', '0',
|
|
'-i', concat_file,
|
|
'-c:v', 'copy', # Copy video codec to preserve alpha
|
|
'-an', # No audio for now
|
|
output_path
|
|
]
|
|
else:
|
|
cmd = [
|
|
'ffmpeg',
|
|
'-y', # Overwrite output
|
|
'-f', 'concat',
|
|
'-safe', '0',
|
|
'-i', concat_file,
|
|
'-c:v', 'copy', # Copy video codec (no re-encoding)
|
|
'-an', # No audio for now
|
|
output_path
|
|
]
|
|
|
|
# Use hardware encoding if requested
|
|
if self.use_nvenc:
|
|
import sys
|
|
if sys.platform == 'darwin':
|
|
encoder = 'hevc_videotoolbox'
|
|
else:
|
|
encoder = 'hevc_nvenc'
|
|
|
|
# Re-encode with hardware acceleration
|
|
cmd = [
|
|
'ffmpeg',
|
|
'-y',
|
|
'-f', 'concat',
|
|
'-safe', '0',
|
|
'-i', concat_file,
|
|
'-c:v', encoder,
|
|
'-preset', 'slow',
|
|
'-b:v', bitrate,
|
|
'-pix_fmt', 'yuv420p',
|
|
'-an',
|
|
output_path
|
|
]
|
|
|
|
logger.info(f"Running concatenation command: {' '.join(cmd)}")
|
|
|
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
|
|
|
if result.returncode != 0:
|
|
logger.error(f"FFmpeg concatenation failed: {result.stderr}")
|
|
return False
|
|
|
|
logger.info(f"Successfully concatenated segments to: {output_path}")
|
|
|
|
# Clean up concat file
|
|
try:
|
|
os.remove(concat_file)
|
|
except:
|
|
pass
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during concatenation: {e}")
|
|
return False
|
|
|
|
def copy_audio_from_original(self, original_video: str, processed_video: str,
|
|
final_output: str) -> bool:
|
|
"""
|
|
Copy audio track from original video to processed video.
|
|
|
|
Args:
|
|
original_video: Path to original video with audio
|
|
processed_video: Path to processed video without audio
|
|
final_output: Path for final output with audio
|
|
|
|
Returns:
|
|
True if successful
|
|
"""
|
|
if not self.preserve_audio:
|
|
logger.info("Audio preservation disabled, skipping audio copy")
|
|
return True
|
|
|
|
try:
|
|
# Check if original video has audio
|
|
probe_cmd = [
|
|
'ffprobe',
|
|
'-v', 'error',
|
|
'-select_streams', 'a:0',
|
|
'-show_entries', 'stream=codec_type',
|
|
'-of', 'csv=p=0',
|
|
original_video
|
|
]
|
|
|
|
result = subprocess.run(probe_cmd, capture_output=True, text=True)
|
|
|
|
if result.returncode != 0 or result.stdout.strip() != 'audio':
|
|
logger.warning("Original video has no audio track")
|
|
# Just copy the processed video
|
|
import shutil
|
|
shutil.copy2(processed_video, final_output)
|
|
return True
|
|
|
|
# Copy audio from original to processed video
|
|
cmd = [
|
|
'ffmpeg',
|
|
'-y',
|
|
'-i', processed_video, # Video input
|
|
'-i', original_video, # Audio input
|
|
'-c:v', 'copy', # Copy video stream
|
|
'-c:a', 'copy', # Copy audio stream
|
|
'-map', '0:v:0', # Map video from first input
|
|
'-map', '1:a:0', # Map audio from second input
|
|
'-shortest', # Match duration to shortest stream
|
|
final_output
|
|
]
|
|
|
|
logger.info("Copying audio from original video...")
|
|
|
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
|
|
|
if result.returncode != 0:
|
|
logger.error(f"FFmpeg audio copy failed: {result.stderr}")
|
|
return False
|
|
|
|
logger.info(f"Successfully added audio to final video: {final_output}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error copying audio: {e}")
|
|
return False
|
|
|
|
def assemble_final_video(self, segments_dir: str, original_video: str,
|
|
output_path: str, bitrate: str = "50M") -> bool:
|
|
"""
|
|
Complete pipeline to assemble final video with audio.
|
|
|
|
Args:
|
|
segments_dir: Directory containing processed segments
|
|
original_video: Path to original video (for audio)
|
|
output_path: Path for final output video
|
|
bitrate: Output video bitrate
|
|
|
|
Returns:
|
|
True if successful
|
|
"""
|
|
logger.info("Starting final video assembly...")
|
|
|
|
# Step 1: Concatenate segments
|
|
temp_concat_path = os.path.join(os.path.dirname(output_path), "temp_concat.mp4")
|
|
|
|
if not self.concatenate_segments(segments_dir, temp_concat_path, bitrate):
|
|
logger.error("Failed to concatenate segments")
|
|
return False
|
|
|
|
# Step 2: Add audio from original
|
|
if self.preserve_audio and file_exists(original_video):
|
|
success = self.copy_audio_from_original(original_video, temp_concat_path, output_path)
|
|
|
|
# Clean up temp file
|
|
try:
|
|
os.remove(temp_concat_path)
|
|
except:
|
|
pass
|
|
|
|
return success
|
|
else:
|
|
# No audio to add, just rename temp file
|
|
import shutil
|
|
try:
|
|
shutil.move(temp_concat_path, output_path)
|
|
logger.info(f"Final video saved to: {output_path}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error moving final video: {e}")
|
|
return False
|
|
|
|
def verify_segment_completeness(self, segments_dir: str) -> tuple[bool, List[int]]:
|
|
"""
|
|
Verify all segments have been processed.
|
|
|
|
Args:
|
|
segments_dir: Directory containing segments
|
|
|
|
Returns:
|
|
Tuple of (all_complete, missing_segments)
|
|
"""
|
|
segments = get_segments_directories(segments_dir)
|
|
missing_segments = []
|
|
|
|
for i, segment in enumerate(segments):
|
|
segment_dir = os.path.join(segments_dir, segment)
|
|
if self.output_mode == "alpha_channel":
|
|
output_video = os.path.join(segment_dir, f"output_{i}.mov")
|
|
else:
|
|
output_video = os.path.join(segment_dir, f"output_{i}.mp4")
|
|
|
|
if not file_exists(output_video):
|
|
missing_segments.append(i)
|
|
|
|
all_complete = len(missing_segments) == 0
|
|
|
|
if all_complete:
|
|
logger.info(f"All {len(segments)} segments have been processed")
|
|
else:
|
|
logger.warning(f"Missing output for segments: {missing_segments}")
|
|
|
|
return all_complete, missing_segments |