1
0
Files
scripts/bin/vcmp
T
2025-01-26 17:12:49 +01:00

359 lines
13 KiB
Python
Executable File

#!/usr/bin/env python3
import os
import pathlib
import subprocess
import shutil
import sys
import argparse
import tempfile
import logging
import time
import psutil
video_extensions: list[str] = ["mp4", "mkv", "avi", "wmv", "mov", "m4v", "ts", "flv", "mpg"]
total_files: int = 0
def is_system_busy() -> bool:
"""Returns true if the cpu is doing something else. maybe."""
total_cpu = 0
our_pids = [os.getpid()]
try:
# add child processes
our_pids.extend(p.pid for p in psutil.Process().children(recursive=True))
except:
pass
for proc in psutil.process_iter(['cpu_percent', 'pid']):
if proc.info['pid'] not in our_pids:
total_cpu += proc.info['cpu_percent']
#print(total_cpu)
cpu_count = psutil.cpu_count()
if cpu_count is None:
cpu_count = 1
return total_cpu > 50 * cpu_count # TODO: Parameter
def acquire_file_lock(filepath: pathlib.Path) -> bool:
lock_file = filepath.with_suffix(filepath.suffix + '.lock')
try:
# Atomic file creation
lock_fd = os.open(str(lock_file), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
with os.fdopen(lock_fd, 'w') as f:
f.write(f"{os.getpid()}\n{time.time()}\n")
return True
except FileExistsError:
try:
lock_age = time.time() - lock_file.stat().st_mtime
if lock_age > 86400: # 24 hours
lock_file.unlink()
return acquire_file_lock(filepath)
except FileNotFoundError:
return acquire_file_lock(filepath)
return False
def release_file_lock(filepath: pathlib.Path) -> None:
lock_file = filepath.with_suffix(filepath.suffix + '.lock')
try:
lock_file.unlink()
except FileNotFoundError:
pass
def setup_logging(verbose: bool) -> None:
"""Configure logging to both file and stdout."""
log_level = logging.DEBUG if verbose else logging.INFO
# Use local log directory
xdg_state_home = os.environ.get('XDG_STATE_HOME', str(pathlib.Path.home() / '.local' / 'state'))
log_dir = pathlib.Path(xdg_state_home) / 'vcmp'
log_dir.mkdir(parents=True, exist_ok=True)
log_file = log_dir / 'vcmp.log'
# Create console handler
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(log_level)
console_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
# Create file handler
file_handler = logging.FileHandler(log_file)
file_handler.setLevel(logging.DEBUG) # Always log everything to file
file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
# Set up root logger
root_logger = logging.getLogger()
root_logger.setLevel(logging.DEBUG)
root_logger.addHandler(console_handler)
root_logger.addHandler(file_handler)
def ffmpeg_command(input: pathlib.Path, output: pathlib.Path, nice: int = 20) -> list[str]:
return [
"nice", "-n", str(nice),
"ffmpeg",
"-progress", "pipe:1", # Send progress info to stdout
"-i", str(input),
"-vf", "fps='if(gte(source_fps,100),source_fps/4,if(gte(source_fps,50),source_fps/2,source_fps))',scale=-2:'min(1080,ih)'",
"-c:v", "libsvtav1",
"-crf", "40",
"-c:a", "aac",
"-b:a", "64k",
str(output)
]
def clean_directory(directory_path: pathlib.Path) -> None:
"""Yeet all files and directories from a directory."""
for item in directory_path.iterdir():
if item.is_dir():
shutil.rmtree(item) # Yeet the directory
else:
item.unlink() # Yeet the file
def is_video(path: pathlib.Path) -> bool:
return path.suffix.lower()[1:] in video_extensions
def get_video_duration(input_file: pathlib.Path) -> float:
cmd = [
"ffprobe",
"-v", "error",
"-show_entries", "format=duration",
"-of", "default=noprint_wrappers=1:nokey=1",
str(input_file)
]
try:
result = subprocess.run(cmd, capture_output=True, text=True, check=True)
return float(result.stdout)
except (subprocess.CalledProcessError, ValueError):
logging.error(f"Error getting duration for {input_file}")
return 1.0
def run_progress(process, duration, file_index=0) -> bool:
current_time = None
last_update = 0
progress = 0
UPDATE_INTERVAL = 1
start_time = time.time()
while True:
line = process.stdout.readline() if process.stdout else None
# Check if we've reached the end of output (no more lines to read)
# and if the process has finished (poll() returns exit code instead of None)
if not line and process.poll() is not None:
break
if line and line.startswith('out_time_ms='):
try:
time_str = line.split('=')[1].strip()
current_time = float(time_str) / 1000000
current_time = min(current_time, duration)
progress = (current_time / duration) * 100
elapsed = time.time() - start_time
speed = current_time / elapsed if elapsed > 0 else 0
if time.time() - last_update > UPDATE_INTERVAL:
if is_system_busy():
return False
total_progress = (file_index / total_files) * 100
sys.stdout.write(
f"\rProcessing ({file_index}/{total_files} [{total_progress:.1f}%]) :"
f"{progress:5.1f}% "
f"({current_time:.1f} / {duration:.1f} seconds) "
f"[{speed:.2f}x]"
f"{'.' * int(progress / 5)}{' ' * (20 - int(progress / 5))}"
)
sys.stdout.flush()
last_update = time.time()
except:
pass
return True
def run(filepath: pathlib.Path, cmd: list[str], temp_path: pathlib.Path, file_index: int) -> None:
duration = get_video_duration(filepath)
xdg_state_home = os.environ.get('XDG_STATE_HOME', str(pathlib.Path.home() / '.local' / 'state'))
ffmpeg_log_path = pathlib.Path(xdg_state_home) / 'vcmp' / 'ffmpeg.log'
ffmpeg_log_path.parent.mkdir(parents=True, exist_ok=True)
with open(ffmpeg_log_path, 'a') as log_file:
process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=log_file,
universal_newlines=True
)
if not run_progress(process, duration, file_index):
process.kill()
release_file_lock(filepath)
if temp_path.exists():
temp_path.unlink()
sys.exit(0)
sys.stdout.write("\n")
sys.stdout.flush()
# Wait for the process to complete and get the return code
process.wait()
if process.returncode != 0:
logging.error(f"FFmpeg exited with code {process.returncode}")
if temp_path.exists():
temp_path.unlink()
release_file_lock(filepath)
raise subprocess.CalledProcessError(process.returncode, cmd)
def process_single_file(
filepath: pathlib.Path,
target_path: pathlib.Path,
temp_dir: pathlib.Path,
remove_source: bool,
nice: int,
file_index: int
) -> bool: # Return success/failure status
if not acquire_file_lock(filepath):
logging.info(f"Skipping {filepath}: already being processed")
return False
temp_path = temp_dir / filepath.with_suffix(".mp4").name
clean_directory(temp_dir) # Should be empty if mktemp was used, but clear it regardless.
logging.info(f"Processing {filepath} to {target_path}")
logging.debug(f"Using temporary file: {temp_path}")
cmd = ffmpeg_command(filepath, temp_path, nice)
logging.debug(f"Executing: {' '.join(cmd)}")
try:
run(filepath, cmd, temp_path, file_index)
except KeyboardInterrupt:
logging.warning("Process interrupted by user")
if temp_path.exists():
temp_path.unlink()
release_file_lock(filepath)
raise # Re-raise KeyboardInterrupt to exit the program
except subprocess.CalledProcessError as e:
logging.error(f"FFmpeg error processing {filepath}: {e}")
if temp_path.exists():
temp_path.unlink()
release_file_lock(filepath)
return False
else:
# File successfully converted to av1, move to destination
target_path.parent.mkdir(parents=True, exist_ok=True)
success = False
try:
# Use cp to copy file since it handles permissions better
subprocess.run(["cp", str(temp_path), str(target_path)], check=True)
temp_path.unlink()
success = True
except subprocess.CalledProcessError as e:
logging.error(f"Failed to copy file: {e}")
if temp_path.exists():
temp_path.unlink()
if success:
logging.info(f"Successfully compressed {filepath}")
if remove_source:
logging.info(f"Removing source file {filepath}")
filepath.unlink()
release_file_lock(filepath)
return True
release_file_lock(filepath)
return False
def process_directory(
*,
source_dir: pathlib.Path,
destination_dir: pathlib.Path,
temp_dir: pathlib.Path,
remove_source: bool = False,
nice: int = 20,
count_only: bool = False
) -> int:
"""Process all video files in source directory recursively."""
count = 0
successful = 0
for dirpath, _dirnames, filenames in os.walk(source_dir):
for filename in filenames:
filepath = pathlib.Path(dirpath) / filename
# File might have been deleted since starting script
if not filepath.exists() or not is_video(filepath):
logging.debug(f"Skipping {filepath}: not a video file or doesn't exist")
continue
if ".cmp." in filepath.name:
logging.debug(f"Skipping {filepath}: already compressed")
continue
# Determine destination path, remove source if already exists and requested
target_path = destination_dir / filepath.relative_to(source_dir)
target_path = target_path.with_suffix(".cmp.mp4")
if target_path.exists():
if remove_source:
logging.info(f"Removing source file {filepath} as compressed version exists")
filepath.unlink()
logging.debug(f"Skipping {filepath}: target already exists")
continue
count += 1
if count_only:
continue
if process_single_file(filepath, target_path, temp_dir, remove_source, nice, count):
successful += 1
if not count_only:
logging.info(f"Processed {count} files, {successful} successful")
return count
def main():
parser = argparse.ArgumentParser(
description="""
Video compression utility that converts videos to AV1 format with reduced quality/size.
Converts to 1080p max resolution, reduces high framerates, and uses low bitrate audio.
The compressed output files will have '.cmp' added to the filename.
"""
)
parser.add_argument('-s', '--source-dir', type=pathlib.Path, default=pathlib.Path.cwd(),
help='Source directory containing videos to compress (default: current directory)')
parser.add_argument('-d', '--destination-dir', type=pathlib.Path,
help='Destination directory for compressed videos (default: same as source)')
parser.add_argument('-t', '--temp-dir', type=pathlib.Path, default=pathlib.Path(tempfile.mkdtemp()),
help='Temporary directory for processing (default: system temp directory)')
parser.add_argument('-r', '--remove-source', action='store_true',
help='Remove source files after successful compression')
parser.add_argument('-v', '--verbose', action='store_true',
help='Enable verbose logging output')
parser.add_argument('-n', '--nice', type=lambda x: int(x) if 0 <= int(x) <= 39 else exec('raise ValueError("Nice value must be between 0 and 39")'),
default=20,
help='Nice value for process priority (default: 20). Valid values are 0 (highest priority) to 39 (lowest priority). Lower values use more system resources.')
parser.add_argument('-i', '--ignore-busy', action='store_true',
help='Ignore system busy check and run anyway')
args = parser.parse_args()
setup_logging(bool(args.verbose))
source_dir = args.source_dir
destination_dir = args.destination_dir if args.destination_dir else args.source_dir
temp_dir = args.temp_dir
if not args.ignore_busy and is_system_busy():
logging.error("System is busy, try again later or use --ignore-busy flag to override")
sys.exit(1)
logging.info(f"Starting video compression from {source_dir} to {destination_dir}")
global total_files
total_files = process_directory(
source_dir=source_dir,
destination_dir=destination_dir,
temp_dir=temp_dir,
count_only=True
)
process_directory(
source_dir=source_dir,
destination_dir=destination_dir,
temp_dir=temp_dir,
remove_source=args.remove_source,
nice=args.nice
)
if __name__ == "__main__":
main()