GigaMux/encode_dvr.py

import os
import subprocess
from pathlib import Path
import json
import logging
from datetime import datetime
import shutil

# ANSI color codes
class Colors:
    PURPLE = '\033[95m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    ENDC = '\033[0m'

def get_gpu_selection():
    while True:
        print(f"\n{Colors.BLUE}Select GPU slot:{Colors.ENDC}")
        print("0 - First GPU")
        print("1 - Second GPU")
        print("2 - Third GPU")
        gpu = input(f"{Colors.YELLOW}Enter GPU number (0-2):{Colors.ENDC} ").strip()
        if gpu in ['0', '1', '2']:
            return gpu
        print(f"{Colors.RED}Invalid selection. Please try again.{Colors.ENDC}")

# Set up logging
log_dir = "logs"
os.makedirs(log_dir, exist_ok=True)
log_file = os.path.join(log_dir, f"encode_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
logging.basicConfig(filename=log_file, level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def get_file_info(input_file):
    cmd = [
        'ffprobe',
        '-v', 'error',
        '-show_entries', 'format=duration,size:stream=codec_name,width,height,r_frame_rate,channels,channel_layout',
        '-of', 'json',
        input_file
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    return json.loads(result.stdout)

def get_audio_labels(input_file):
    cmd = [
        'ffprobe',
        '-v', 'error',
        '-select_streams', 'a',
        '-show_entries', 'stream=index:stream_tags=title',
        '-of', 'json',
        input_file
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    info = json.loads(result.stdout)
    labels = []
    for stream in info.get('streams', []):
        title = stream.get('tags', {}).get('title', None)
        labels.append(title)
    return labels

def format_size(size_bytes):
    for unit in ['B', 'KB', 'MB', 'GB']:
        if size_bytes < 1024:
            return f"{size_bytes:.2f} {unit}"
        size_bytes /= 1024
    return f"{size_bytes:.2f} TB"

def encode_dvr(input_file, output_dir, gpu):
    input_path = Path(input_file)
    output_path = Path(output_dir) / f"{input_path.stem}{input_path.suffix}"

    # Get file info for logging
    file_info = get_file_info(str(input_path))
    input_size = int(file_info['format']['size'])
    duration = float(file_info['format']['duration'])
    
    logging.info(f"Processing file: {input_path}")
    logging.info(f"Input size: {format_size(input_size)}")
    logging.info(f"Duration: {duration:.2f} seconds")
    
    print(f"\n{Colors.BLUE}Processing file: {input_path}{Colors.ENDC}")
    print(f"Input size: {format_size(input_size)}")
    print(f"Duration: {duration:.2f} seconds")
    
    # Log stream information
    for i, stream in enumerate(file_info.get('streams', [])):
        stream_type = 'Video' if stream.get('codec_name', '').startswith('h') else 'Audio'
        logging.info(f"Stream {i} ({stream_type}):")
        for key, value in stream.items():
            logging.info(f"  {key}: {value}")

    # Skip if output file already exists
    if output_path.exists():
        output_size = output_path.stat().st_size
        logging.info(f"Skipping {input_path} - output already exists: {output_path}")
        logging.info(f"Output size: {format_size(output_size)}")
        print(f"{Colors.YELLOW}Skipping {input_path} - output already exists{Colors.ENDC}")
        return

    # Get audio labels
    audio_labels = get_audio_labels(str(input_path))
    logging.info(f"Audio labels: {audio_labels}")

    # FFmpeg command with NVIDIA HEVC encoder and maximum quality
    cmd = [
        'ffmpeg',
        '-v', 'verbose',
        '-progress', 'pipe:1',  # Force progress output to stdout
        '-stats',              # Show encoding statistics
        '-stats_period', '0.5', # Update stats every 0.5 seconds
        '-i', str(input_path),
        '-c:v', 'hevc_nvenc',
        '-gpu', gpu,
        '-preset', 'p7',
        '-tune', 'hq',
        '-rc', 'vbr',  # Variable bitrate mode
        '-rc-lookahead', '32',  # Increase lookahead for better VBR decisions
        '-spatial-aq', '1',  # Enable spatial AQ for better quality
        '-aq-strength', '15',  # Maximum allowed AQ strength
        '-cq', '35',
        '-b:v', '0',  # Let VBR determine bitrate
        '-c:a', 'copy',
        '-map', '0',
    ]
    # Add metadata for each audio stream if label exists
    for idx, label in enumerate(audio_labels):
        if label:
            cmd += [f'-metadata:s:a:{idx}', f'title={label}']
    cmd.append(str(output_path))

    try:
        # Run FFmpeg with timeout
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True
        )
        
        # Set timeout to 5 minutes
        timeout = 300
        start_time = datetime.now()
        
        # Log FFmpeg output in real-time
        while True:
            # Check for timeout
            if (datetime.now() - start_time).total_seconds() > timeout:
                process.kill()
                raise TimeoutError(f"FFmpeg process timed out after {timeout} seconds")
                
            output = process.stderr.readline()
            if output == '' and process.poll() is not None:
                break
            if output:
                logging.info(f"FFmpeg: {output.strip()}")
                print(f"{Colors.PURPLE}FFmpeg: {output.strip()}{Colors.ENDC}")

        if process.returncode == 0:
            # Get output file info
            output_info = get_file_info(str(output_path))
            output_size = int(output_info['format']['size'])
            compression_ratio = input_size / output_size if output_size > 0 else 0
            
            logging.info(f"Successfully encoded: {output_path}")
            logging.info(f"Output size: {format_size(output_size)}")
            logging.info(f"Compression ratio: {compression_ratio:.2f}x")
            print(f"{Colors.GREEN}Successfully encoded: {output_path}{Colors.ENDC}")
            print(f"Compression ratio: {compression_ratio:.2f}x")
        else:
            logging.error(f"FFmpeg process failed with return code {process.returncode}")
            print(f"{Colors.RED}FFmpeg process failed with return code {process.returncode}{Colors.ENDC}")
            
    except TimeoutError as e:
        logging.error(f"Timeout error: {e}")
        print(f"{Colors.RED}Timeout error: {e}{Colors.ENDC}")
    except subprocess.CalledProcessError as e:
        logging.error(f"Error encoding {input_path}: {e}")
        print(f"{Colors.RED}Error encoding {input_path}: {e}{Colors.ENDC}")
    except Exception as e:
        logging.error(f"Unexpected error: {e}")
        print(f"{Colors.RED}Unexpected error: {e}{Colors.ENDC}")

if __name__ == "__main__":
    # Get GPU selection
    gpu = get_gpu_selection()
    logging.info(f"Selected GPU: {gpu}")
    
    input_dir = "input"
    output_dir = "output"
    os.makedirs(output_dir, exist_ok=True)
    
    # Get list of files to process
    files = [f for f in os.listdir(input_dir) if f.endswith(('.mp4', '.DVR.mp4'))]
    total_files = len(files)
    
    if total_files == 0:
        logging.info("No files to process in input directory")
        print(f"{Colors.YELLOW}No files to process in input directory{Colors.ENDC}")
    else:
        logging.info(f"Found {total_files} files to process")
        print(f"{Colors.BLUE}Found {total_files} files to process{Colors.ENDC}")
        
        for i, file in enumerate(files, 1):
            input_file = os.path.join(input_dir, file)
            logging.info(f"Processing file {i}/{total_files}: {file}")
            print(f"\n{Colors.BLUE}Processing file {i}/{total_files}: {file}{Colors.ENDC}")
            encode_dvr(input_file, output_dir, gpu)