From 25e8c0293e4b8bc52984c6c6a1711f98594603d9 Mon Sep 17 00:00:00 2001 From: Maxim Vershinin Date: Wed, 2 Sep 2026 11:42:47 +0500 Subject: [PATCH] feat (core): allows user to add an external cover image to split tracks --- audio_splitter/core.py | 43 +++++++++++++-------- audio_splitter/ffmpeg.py | 81 ++++++++++++++++++++++++++++++++++++++++ audio_splitter/main.py | 26 +++++++++++++ 3 files changed, 135 insertions(+), 15 deletions(-) diff --git a/audio_splitter/core.py b/audio_splitter/core.py index 6aa9865..3bc3f34 100644 --- a/audio_splitter/core.py +++ b/audio_splitter/core.py @@ -59,7 +59,7 @@ def split_audio(input_file: str, output_directory: str, tracks: List[Dict[str, A - merge_comments: bool - comment_separator: str - delete_original: bool - - cover_image: str or None (optional, set by web backend) + - cover_image: List[str] or None (optional, CLI list of cover image paths) Raises: RuntimeError: If no audio stream is found. @@ -152,24 +152,33 @@ def split_audio(input_file: str, output_directory: str, tracks: List[Dict[str, A # -------------------------------------------------------------------------- # 7. Handle attached picture (cover art) # -------------------------------------------------------------------------- - cover_image_path = getattr(args, 'cover_image', None) - if cover_image_path and not os.path.exists(cover_image_path): - cover_image_path = None - - if not cover_image_path and not args.drop_video and stream_info.get('has_video'): + user_cover_images = getattr(args, 'cover_image', None) + # Resolve per-track cover images: single image reused, or one-per-track + track_cover_images = [] + if user_cover_images: + for img in user_cover_images: + if os.path.exists(img): + track_cover_images.append(img) + else: + print(f"Warning: Cover image not found, skipping: {img}") + # Extract cover from input if no user cover and input has video + if not track_cover_images and not args.drop_video and stream_info.get('has_video'): if is_attached_picture(input_file): cover_image_path = os.path.join(output_directory, 'cover.png') if extract_cover_image(input_file, cover_image_path): + track_cover_images.append(cover_image_path) print("Extracted cover image for all tracks.") else: - cover_image_path = None + track_cover_images = [] # Check if we need special handling (e.g., Opus files need opustags) + # Use output_audio_codec (not input_audio_codec) so the handler is resolved + # against the actual output codec (e.g., transcoding aac→opus in ogg). input_audio_codec = stream_info.get('audio_codec') - cover_handler = get_handler(output_container, input_audio_codec) - needs_drop = needs_drop_video(output_container, input_audio_codec) - if cover_image_path and needs_drop: - print(f"Using special handler for {output_container} + {input_audio_codec}") + cover_handler = get_handler(output_container, output_audio_codec) + needs_drop = needs_drop_video(output_container, output_audio_codec) + if track_cover_images and needs_drop: + print(f"Using special handler for {output_container} + {output_audio_codec}") print("Temporarily dropping video for opustags post-processing.") # -------------------------------------------------------------------------- @@ -219,10 +228,14 @@ def split_audio(input_file: str, output_directory: str, tracks: List[Dict[str, A # Then build command with these mapped encoders # Use temporary drop_video for handlers that need it - effective_drop_video = args.drop_video or (cover_image_path and needs_drop) + effective_drop_video = args.drop_video or (track_cover_images and needs_drop) # For handlers that need opustags post-processing, don't pass cover_image_path # to ffmpeg - it will process audio-only, then handler adds cover afterward - ffmpeg_cover_path = None if (cover_image_path and needs_drop) else cover_image_path + per_track_cover = ( + track_cover_images[idx - 1] if len(track_cover_images) == len(tracks) + else (track_cover_images[0] if track_cover_images else None) + ) + ffmpeg_cover_path = None if (per_track_cover and needs_drop) else per_track_cover cmd = build_ffmpeg_command( input_file=input_file, start_seconds=start_seconds, @@ -251,11 +264,11 @@ def split_audio(input_file: str, output_directory: str, tracks: List[Dict[str, A else: print(f" -> Saved to: {output_path}") # Apply cover image handler if needed - if cover_image_path and not args.drop_video and needs_drop: + if per_track_cover and needs_drop: print(f"Applying cover image via {output_container} handler...") if not cover_handler( input_file=input_file, - cover_image_path=cover_image_path, + cover_image_path=per_track_cover, output_path=output_path, stream_info=stream_info, audio_codec=audio_enc, diff --git a/audio_splitter/ffmpeg.py b/audio_splitter/ffmpeg.py index 8d4cb35..741bbd0 100644 --- a/audio_splitter/ffmpeg.py +++ b/audio_splitter/ffmpeg.py @@ -2,6 +2,7 @@ """FFmpeg/FFprobe interaction utilities.""" import json +import os import subprocess from typing import Dict, List, Optional, Tuple @@ -201,6 +202,86 @@ def get_container_format(input_file: str) -> Optional[str]: return mapping.get(format_name, format_name) +MAX_COVER_IMAGE_SIZE_MB = 4 + + +def has_cover_or_video(input_file: str) -> bool: + """ + Check if the input file contains a cover image (attached picture) or any video track. + + Returns: + True if a cover image or video stream is detected, False otherwise. + """ + cmd = [ + 'ffprobe', '-v', 'quiet', + '-print_format', 'json', + '-select_streams', 'v', + '-show_entries', 'stream=codec_type,codec_name,disposition', + input_file + ] + result = subprocess.run(cmd, capture_output=True, text=True, check=False) + if result.returncode != 0: + return False + + try: + data = json.loads(result.stdout) + streams = data.get('streams', []) + return len(streams) > 0 + except (json.JSONDecodeError, KeyError): + return False + + +def validate_cover_images(cover_images: List[str], num_tracks: int) -> None: + """ + Validate cover image(s) before splitting. + + Checks: + - Each file exists + - Only PNG and JPEG formats are accepted + - Single image OR count matches number of tracks + - File size warning for images > 4 MB + + Args: + cover_images: List of cover image paths (single image or multiple). + num_tracks: Number of tracks in the tracklist. + + Raises: + ValueError: If validation fails. + """ + if len(cover_images) == 0: + return + + if len(cover_images) == 1 and num_tracks > 1: + # Single image is fine - will be reused for all tracks + pass + elif len(cover_images) == num_tracks: + pass + else: + raise ValueError( + f"Cover image count ({len(cover_images)}) must be 1 or match " + f"the number of tracks ({num_tracks})." + ) + + for img_path in cover_images: + if not os.path.exists(img_path): + raise ValueError(f"Cover image not found: {img_path}") + + ext = os.path.splitext(img_path)[1].lower() + if ext not in ('.png', '.jpg', '.jpeg'): + raise ValueError( + f"Unsupported cover image format '{ext}' for '{img_path}'. " + f"Only PNG and JPEG are supported." + ) + + size_mb = os.path.getsize(img_path) / (1024 * 1024) + if size_mb > MAX_COVER_IMAGE_SIZE_MB: + print( + f"Warning: Cover image '{img_path}' is {size_mb:.1f} MB " + f"(exceeds {MAX_COVER_IMAGE_SIZE_MB} MB limit). " + f"Large images may cause issues." + ) + + def is_attached_picture(input_file: str) -> bool: """ Check if the input file has a video stream that is an attached picture (cover art). diff --git a/audio_splitter/main.py b/audio_splitter/main.py index fce53e3..38e9020 100644 --- a/audio_splitter/main.py +++ b/audio_splitter/main.py @@ -34,6 +34,7 @@ from .defaults import ( ) from .core import split_audio from .tracklist import read_tracklist, parse_format +from .ffmpeg import has_cover_or_video, validate_cover_images def main(): @@ -156,6 +157,16 @@ def main(): help="Tracklist format (default: %(default)s)" ) + # Cover image + parser.add_argument( + '--cover-image', + nargs='+', + metavar='IMAGE', + default=None, + help="Cover image path(s). Single image applied to all tracks, " + "or one per track (must match track count)." + ) + # Other parser.add_argument( '--delete-original', @@ -208,6 +219,21 @@ def main(): print(f"Found {len(tracks)} tracks.") + # Validate cover images if provided + if args.cover_image: + # Check if input already has cover image or video track + # Allow if --drop-video is specified (user wants to discard existing video) + if has_cover_or_video(args.input_file) and not args.drop_video: + print("Error: Input file already contains a cover image or video track. " + "Remove it first, or use --drop-video to discard video streams.") + sys.exit(1) + + try: + validate_cover_images(args.cover_image, len(tracks)) + except ValueError as e: + print(f"Error: {e}") + sys.exit(1) + if args.dry_run: print("\nParsed tracklist:") print("-" * 60)