video_echo_pipeline.py

#!/usr/bin/env python3
"""
Video Echo Pipeline
Creates 300 progressively echo-enhanced images and stitches them into a 60fps video.
"""

import os
import sys
import subprocess
import time
import numpy as np
from pathlib import Path
from PIL import Image

class VideoEchoProcessor:
    def __init__(self, input_dir="input", output_dir="output", video_dir="video"):
        self.input_dir = Path(input_dir)
        self.output_dir = Path(output_dir)
        self.video_dir = Path(video_dir)
        self.bitmap_dir = Path("input-bitmap")
        self.raw_dir = Path("output-raw")
        self.audio_dir = Path("temp-audio")
        
        # Create all necessary directories
        for dir_path in [self.output_dir, self.video_dir, self.bitmap_dir, self.raw_dir, self.audio_dir]:
            dir_path.mkdir(exist_ok=True)
    
    def convert_to_bitmap(self):
        """Convert JPG/PNG files to bitmap format"""
        print("Step 1: Converting images to bitmap format...")
        
        supported_formats = {'.jpg', '.jpeg', '.png'}
        image_files = []
        
        for file_path in self.input_dir.iterdir():
            if file_path.is_file() and file_path.suffix.lower() in supported_formats:
                image_files.append(file_path)
        
        if not image_files:
            print(f"No JPG or PNG files found in '{self.input_dir}' directory.")
            return False
        
        print(f"Found {len(image_files)} image file(s) to convert:")
        
        for i, file_path in enumerate(image_files, 1):
            try:
                print(f"[{i}/{len(image_files)}] Converting {file_path.name}...")
                
                with Image.open(file_path) as img:
                    if img.mode != 'RGB':
                        img = img.convert('RGB')
                    
                    output_filename = file_path.stem + '.bmp'
                    output_path = self.bitmap_dir / output_filename
                    
                    img.save(output_path, 'BMP')
                    print(f"  ✓ Saved as {output_filename}")
                    
            except Exception as e:
                print(f"  ✗ Error converting {file_path.name}: {e}")
                return False
        
        return True
    
    def bitmap_to_raw_audio(self):
        """Convert bitmap files to raw audio format for processing"""
        print("Step 2: Converting bitmaps to raw audio...")
        
        bitmap_files = list(self.bitmap_dir.glob("*.bmp"))
        
        if not bitmap_files:
            print(f"No bitmap files found in '{self.bitmap_dir}' directory.")
            return False
        
        print(f"Found {len(bitmap_files)} bitmap file(s) to convert:")
        
        for i, bitmap_file in enumerate(bitmap_files, 1):
            try:
                print(f"[{i}/{len(bitmap_files)}] Converting {bitmap_file.name}...")
                
                # Read bitmap file
                with open(bitmap_file, 'rb') as f:
                    bitmap_data = f.read()
                
                # Skip the header (first 255 bytes as per the tutorial)
                audio_data = bitmap_data[255:]
                
                # Convert to 8-bit unsigned audio
                audio_array = np.frombuffer(audio_data, dtype=np.uint8)
                
                # Normalize to 16-bit signed audio for better processing
                audio_16bit = ((audio_array.astype(np.float32) - 128) * 256).astype(np.int16)
                
                # Save as raw audio file
                audio_filename = bitmap_file.stem + '.raw'
                audio_path = self.audio_dir / audio_filename
                
                with open(audio_path, 'wb') as f:
                    f.write(audio_16bit.tobytes())
                
                print(f"  ✓ Converted to {audio_filename}")
                
            except Exception as e:
                print(f"  ✗ Error converting {bitmap_file.name}: {e}")
                return False
        
        return True
    
    def apply_video_progressive_echo(self, num_iterations=300):
        """Apply progressive echo effects for video creation"""
        print(f"Step 3: Applying progressive echo effects ({num_iterations} iterations)...")
        
        # Check if SoX is available
        try:
            subprocess.run(["sox", "--version"], capture_output=True, check=True)
        except FileNotFoundError:
            print("SoX is not installed. Installing...")
            try:
                subprocess.run(["brew", "install", "sox"], check=True)
                print("SoX installed successfully!")
            except subprocess.CalledProcessError:
                print("Failed to install SoX. Please install manually:")
                print("  macOS: brew install sox")
                print("  Linux: sudo apt-get install sox")
                return False
        
        audio_files = list(self.audio_dir.glob("*.raw"))
        
        if not audio_files:
            print(f"No audio files found in '{self.audio_dir}' directory.")
            return False
        
        print(f"Found {len(audio_files)} audio file(s) to process:")
        
        # Clear output directory
        for file in self.raw_dir.glob("*"):
            file.unlink()
        
        for audio_file in audio_files:
            print(f"Processing {audio_file.name} with progressive echo for video...")
            
            # Progressive echo parameters for smooth video
            base_gain_in = 0.98
            base_gain_out = 0.98
            base_delay = 0.05
            base_decay = 0.005
            
            for i in range(num_iterations):
                try:
                    # Calculate progressive values with smooth curves
                    progress = i / (num_iterations - 1)  # 0 to 1
                    
                    # Smooth progression for video
                    gain_in = base_gain_in - (0.5 * (progress ** 1.2))
                    gain_out = base_gain_out - (0.6 * (progress ** 1.0))
                    delay = base_delay + (1.0 - base_delay) * (progress ** 1.5)
                    decay = base_decay + (0.95 - base_decay) * (progress ** 1.8)
                    
                    # Create output filename with iteration number
                    output_file = self.raw_dir / f"{audio_file.stem}_video_echo_{i:04d}.raw"
                    
                    # Build SoX command
                    cmd = ["sox", "-r", "44100", "-e", "signed", "-b", "16", "-c", "1", 
                           str(audio_file), str(output_file),
                           "echo",
                           str(gain_in),
                           str(gain_out),
                           str(delay),
                           str(decay)]
                    
                    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
                    
                    if i % 50 == 0:  # Progress indicator every 50 iterations
                        print(f"  ✓ Iteration {i+1}/{num_iterations} (delay: {delay:.3f}s, decay: {decay:.3f})")
                    
                except subprocess.CalledProcessError as e:
                    print(f"  ✗ Error in iteration {i+1}: {e.stderr}")
                    return False
        
        print(f"  ✓ Completed {num_iterations} progressive echo iterations")
        return True
    
    def fix_bitmap_headers(self):
        """Fix bitmap headers by copying from original to raw files"""
        print("Step 4: Fixing bitmap headers...")
        
        raw_files = list(self.raw_dir.glob("*.raw"))
        
        if not raw_files:
            print(f"No .raw files found in '{self.raw_dir}' directory.")
            return False
        
        print(f"Found {len(raw_files)} raw file(s) to fix:")
        
        # Clear output directory
        for file in self.output_dir.glob("*"):
            file.unlink()
        
        for i, raw_file in enumerate(raw_files, 1):
            try:
                if i % 50 == 0:  # Progress indicator every 50 files
                    print(f"[{i}/{len(raw_files)}] Processing {raw_file.name}...")
                
                # Find corresponding original bitmap file
                base_name = raw_file.stem.split('_video_echo_')[0]
                original_filename = base_name + '.bmp'
                original_file = self.bitmap_dir / original_filename
                
                if not original_file.exists():
                    print(f"  ✗ Original bitmap file '{original_filename}' not found")
                    continue
                
                # Read the first 255 bytes from the original bitmap
                with open(original_file, 'rb') as f:
                    header = f.read(255)  # 0xFF bytes
                
                # Read the raw file content
                with open(raw_file, 'rb') as f:
                    raw_content = f.read()
                
                # Convert back to 8-bit format for bitmap
                raw_array = np.frombuffer(raw_content, dtype=np.int16)
                bitmap_array = ((raw_array.astype(np.float32) / 256) + 128).astype(np.uint8)
                bitmap_content = bitmap_array.tobytes()
                
                # Combine header with processed content
                fixed_content = header + bitmap_content
                
                # Create output filename
                output_filename = raw_file.stem + '.bmp'
                output_file = self.output_dir / output_filename
                
                # Write the fixed bitmap
                with open(output_file, 'wb') as f:
                    f.write(fixed_content)
                
                if i % 50 == 0:  # Progress indicator every 50 files
                    print(f"  ✓ Fixed bitmap saved as {output_filename}")
                
            except Exception as e:
                print(f"  ✗ Error processing {raw_file.name}: {e}")
                return False
        
        print(f"  ✓ Completed fixing headers for {len(raw_files)} files")
        return True
    
    def create_video(self, fps=60):
        """Create video from the processed images using FFmpeg"""
        print(f"Step 5: Creating {fps}fps video from processed images...")
        
        # Check if FFmpeg is available
        try:
            subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
        except FileNotFoundError:
            print("FFmpeg is not installed. Installing...")
            try:
                subprocess.run(["brew", "install", "ffmpeg"], check=True)
                print("FFmpeg installed successfully!")
            except subprocess.CalledProcessError:
                print("Failed to install FFmpeg. Please install manually:")
                print("  macOS: brew install ffmpeg")
                print("  Linux: sudo apt-get install ffmpeg")
                return False
        
        # Get all processed images
        image_files = sorted(list(self.output_dir.glob("*.bmp")))
        
        if not image_files:
            print(f"No processed images found in '{self.output_dir}' directory.")
            return False
        
        print(f"Found {len(image_files)} images to create video from...")
        
        # Create video filename
        base_name = image_files[0].stem.split('_video_echo_')[0]
        video_filename = f"{base_name}_echo_progression_{fps}fps.mp4"
        video_path = self.video_dir / video_filename
        
        # Build FFmpeg command
        # Use pattern matching for input files
        input_pattern = str(self.output_dir / f"{base_name}_video_echo_%04d.bmp")
        
        cmd = [
            "ffmpeg",
            "-y",  # Overwrite output file
            "-framerate", str(fps),
            "-i", input_pattern,
            "-c:v", "libx264",
            "-preset", "medium",
            "-crf", "18",  # High quality
            "-pix_fmt", "yuv420p",
            "-movflags", "+faststart",
            str(video_path)
        ]
        
        print(f"Creating video: {video_filename}")
        print(f"Command: {' '.join(cmd)}")
        
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            print(f"  ✓ Video created successfully: {video_filename}")
            print(f"  ✓ Video saved to: {video_path}")
            
            # Get video info
            info_cmd = ["ffprobe", "-v", "quiet", "-print_format", "json", 
                       "-show_format", "-show_streams", str(video_path)]
            info_result = subprocess.run(info_cmd, capture_output=True, text=True, check=True)
            
            # Extract duration
            import json
            info = json.loads(info_result.stdout)
            duration = float(info['format']['duration'])
            print(f"  ✓ Video duration: {duration:.2f} seconds")
            print(f"  ✓ Frame count: {len(image_files)} frames")
            
            return True
            
        except subprocess.CalledProcessError as e:
            print(f"  ✗ Error creating video: {e.stderr}")
            return False
    
    def run_video_pipeline(self, num_iterations=300, fps=60):
        """Run the complete video echo pipeline"""
        print("Starting video echo pipeline...")
        print("=" * 80)
        print(f"Will create {num_iterations} progressively echo-enhanced images")
        print(f"Then stitch them into a {fps}fps video")
        print("=" * 80)
        
        # Step 1: Convert to bitmap
        if not self.convert_to_bitmap():
            print("Pipeline failed at step 1")
            return False
        
        print()
        
        # Step 2: Convert bitmap to raw audio
        if not self.bitmap_to_raw_audio():
            print("Pipeline failed at step 2")
            return False
        
        print()
        
        # Step 3: Apply progressive effects
        if not self.apply_video_progressive_echo(num_iterations):
            print("Pipeline failed at step 3")
            return False
        
        print()
        
        # Step 4: Fix headers
        if not self.fix_bitmap_headers():
            print("Pipeline failed at step 4")
            return False
        
        print()
        
        # Step 5: Create video
        if not self.create_video(fps):
            print("Pipeline failed at step 5")
            return False
        
        print()
        print("=" * 80)
        print("Video echo pipeline completed successfully!")
        print(f"Created {num_iterations} progressively echo-enhanced images")
        print(f"Created {fps}fps video in: {self.video_dir}")
        print("=" * 80)
        
        return True

def main():
    """Main function"""
    
    # Default parameters
    num_iterations = 300
    fps = 60
    
    # Allow command line arguments
    if len(sys.argv) > 1:
        try:
            num_iterations = int(sys.argv[1])
            if num_iterations < 10 or num_iterations > 1000:
                print("Number of iterations should be between 10 and 1000")
                return
        except ValueError:
            print("Please provide a valid number for iterations")
            return
    
    if len(sys.argv) > 2:
        try:
            fps = int(sys.argv[2])
            if fps < 1 or fps > 120:
                print("FPS should be between 1 and 120")
                return
        except ValueError:
            print("Please provide a valid number for FPS")
            return
    
    # Create processor and run pipeline
    processor = VideoEchoProcessor()
    success = processor.run_video_pipeline(num_iterations, fps)
    
    if success:
        print(f"\nUsage examples:")
        print(f"  python video_echo_pipeline.py                    # Default: 300 iterations, 60fps")
        print(f"  python video_echo_pipeline.py 200               # 200 iterations, 60fps")
        print(f"  python video_echo_pipeline.py 300 30            # 300 iterations, 30fps")
        print(f"  python video_echo_pipeline.py 500 120           # 500 iterations, 120fps")
        print(f"\nVideo features:")
        print(f"  - {num_iterations} frames showing echo progression")
        print(f"  - {fps}fps smooth playback")
        print(f"  - High quality H.264 encoding")
        print(f"  - Duration: ~{num_iterations/fps:.1f} seconds")
    else:
        print("Pipeline failed. Check the error messages above.")

if __name__ == "__main__":
    main()