sbs working phase 1
This commit is contained in:
306
core/video_assembler.py
Normal file
306
core/video_assembler.py
Normal file
@@ -0,0 +1,306 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user