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', '-i', str(input_path), '-c:v', 'hevc_nvenc', '-gpu', gpu, '-preset', 'p7', '-tune', 'hq', '-rc', 'vbr', '-rc-lookahead', '32', '-spatial-aq', '1', '-aq-strength', '15', '-cq', '26', '-b:v', '2500k', '-maxrate', '2500k', '-bufsize', '5000k', '-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 and capture output process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True ) # Log FFmpeg output in real-time while True: 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 subprocess.CalledProcessError as e: logging.error(f"Error encoding {input_path}: {e}") print(f"{Colors.RED}Error encoding {input_path}: {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)