init
This commit is contained in:
194
encode_vtubrec5k.py
Normal file
194
encode_vtubrec5k.py
Normal file
@@ -0,0 +1,194 @@
|
||||
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', '23',
|
||||
'-b:v', '5000k',
|
||||
'-maxrate', '5000k',
|
||||
'-bufsize', '10000k',
|
||||
'-b:a', '192k',
|
||||
'-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)
|
||||
Reference in New Issue
Block a user