Wan2GP/extract_source_images.py

#!/usr/bin/env python3
"""
Extract Source Images from Video Files

This utility extracts source images that were embedded in video files
generated by Wan2GP with source image embedding enabled.
Supports MKV (attachments) and MP4 (cover art) formats.
"""

import sys
import os
import argparse
import glob
from shared.utils.audio_video import extract_source_images


def main():
    parser = argparse.ArgumentParser(
        description="Extract source images from video files with embedded images (MKV attachments or MP4 cover art)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__
    )
    
    parser.add_argument(
        'video_files', 
        nargs='+', 
        help='Video file(s) to extract source images from (MKV or MP4)'
    )
    
    parser.add_argument(
        '-o', '--output-dir',
        help='Output directory for extracted images (default: same as video directory)'
    )
    
    parser.add_argument(
        '-v', '--verbose',
        action='store_true',
        help='Enable verbose output'
    )
    
    args = parser.parse_args()
    
    # Expand glob patterns
    video_files = []
    for pattern in args.video_files:
        if '*' in pattern or '?' in pattern:
            video_files.extend(glob.glob(pattern))
        else:
            video_files.append(pattern)
    
    if not video_files:
        print("No video files found matching the specified patterns.")
        return 1
    
    total_extracted = 0
    
    for video_file in video_files:
        if not os.path.exists(video_file):
            print(f"Warning: File not found: {video_file}")
            continue
            
        if not (video_file.lower().endswith('.mkv') or video_file.lower().endswith('.mp4')):
            print(f"Warning: Skipping unsupported file: {video_file} (only MKV and MP4 supported)")
            continue
        
        if args.verbose:
            print(f"\nProcessing: {video_file}")
        
        # Determine output directory
        if args.output_dir:
            output_dir = args.output_dir
        else:
            # Create subdirectory next to video file
            video_dir = os.path.dirname(video_file) or '.'
            video_name = os.path.splitext(os.path.basename(video_file))[0]
            output_dir = os.path.join(video_dir, f"{video_name}_sources")
        
        try:
            extracted_files = extract_source_images(video_file, output_dir)
            
            if extracted_files:
                print(f"✓ Extracted {len(extracted_files)} source image(s) from {video_file}")
                if args.verbose:
                    for img_file in extracted_files:
                        print(f"  → {img_file}")
                total_extracted += len(extracted_files)
            else:
                print(f"ℹ No source images found in {video_file}")
                
        except Exception as e:
            print(f"✗ Error processing {video_file}: {e}")
    
    print(f"\nTotal: Extracted {total_extracted} source image(s) from {len(video_files)} video file(s)")
    return 0


if __name__ == "__main__":
    sys.exit(main())