batch.py - master - domahes/audio_visualizer_improved

batch.py
322 строки15 KB
Евгений Харченко
Refactor visualizer: split GUI, CLI, batch, and pipeline modules into separate files and add package installation support
02 июн 2026, 11:09
02 июн 2026, 11:0975419c0
"""
Batch processing and merging logic for the audio visualizer.
Phase 4: extracted from main (render_batch, _render_one_for_batch, slate, chapters, tracklist).
Depends on pipeline for per-track render_video.
"""

from __future__ import annotations

import os
import subprocess
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from multiprocessing import cpu_count
from typing import List, Dict, Tuple, Optional

from availability import _SESSION_TMPDIR, RenderCancelled, cv2, clamp
from models import RenderSettings
from utils import _escape_ffconcat_path, _format_hms, _extract_track_metadata, _normalize_tag_value, _build_track_metadata_text
from primitives import _get_bg, _put_text_pil, _get_text_size_pil
from pipeline import render_video


def _ffmetadata_for_chapters(titles: List[str], durations_sec: List[float], gap_duration: float = 0.0) -> str:
    """Create FFMETADATA chapters, accounting for gaps between tracks."""
    if len(titles) != len(durations_sec) or len(titles) == 0:
        raise ValueError("titles and durations_sec must have the same non-zero length")
    gap_duration = max(0.0, float(gap_duration))
    lines = [";FFMETADATA1"]
    time_ms = 0
    for i, (title, dur) in enumerate(zip(titles, durations_sec)):
        start = int(round(time_ms))
        end = int(round(time_ms + max(0.0, dur) * 1000)) - 1
        lines.append("[CHAPTER]")
        lines.append("TIMEBASE=1/1000")
        lines.append(f"START={start}")
        lines.append(f"END={max(start, end)}")
        safe_title = str(title).replace('\n', ' ').strip()
        lines.append(f"title={safe_title}")
        
        # FIX: Add track duration
        time_ms += max(0.0, dur) * 1000
        
        # FIX: Add gap duration if it's not the last track
        if i < len(titles) - 1:
            time_ms += gap_duration * 1000
            
    return "\n".join(lines) + "\n"


def _generate_tracklist_txt(titles: List[str], durations_sec: List[float], gap_duration: float = 0.0, output_path: str = "tracklist.txt") -> str:
    """Генерирует список треков в формате: <Название трека> <Время начала>"""
    if len(titles) != len(durations_sec) or len(titles) == 0:
        raise ValueError("titles and durations_sec must have the same non-zero length")
    gap_duration = max(0.0, float(gap_duration))
    lines = []
    time_sec = 0.0
    
    for i, (title, dur) in enumerate(zip(titles, durations_sec)):
        timestamp = _format_hms(time_sec) # Reuse existing helper (H:MM:SS or MM:SS)
        
        safe_title = str(title).replace('\n', ' ').strip()
        lines.append(f"{safe_title} {timestamp}")
        
        time_sec += max(0.0, float(dur))
        if i < len(titles) - 1:
            time_sec += gap_duration
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write("\n".join(lines))
    
    return output_path


def _generate_slate_video(width: int, height: int, fps: int, duration: float, label: str, 
                          background: str = 'black', font_path: Optional[str] = None, 
                          font_size: int = 40) -> str:
    if duration <= 0:
        raise ValueError("Slate duration must be > 0")
    frames = max(1, int(round(duration * max(1, fps))))
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    path = os.path.join(_SESSION_TMPDIR, f"slate_{abs(hash((label, frames, width, height, fps)))}_tmp_{os.getpid()}.mp4")
    out = cv2.VideoWriter(path, fourcc, max(1, fps), (width, height))
    if not out.isOpened():
        raise RuntimeError("Failed to open VideoWriter for slate")
    
    tw, th = _get_text_size_pil(label, font_path, font_size)
    x = (width - tw) // 2
    y = (height - th) // 2
    
    for i in range(frames):
        frame = _get_bg(width, height, background, i)
        fade_in = clamp(i / 12.0, 0, 1)
        fade_out = clamp((frames - 1 - i) / 12.0, 0, 1)
        alpha = min(fade_in, fade_out)
        
        overlay = frame.copy()
        overlay = _put_text_pil(overlay, label, (x, y), font_path, font_size,
                                color=(240, 240, 240), shadow=True, shadow_offset=2)
        frame = cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0)
        out.write(frame)
    out.release()
    return path


def _render_one_for_batch(ap: str, base_settings: RenderSettings) -> Tuple[str, str, float]:
    file_title = os.path.splitext(os.path.basename(ap))[0]
    meta = _extract_track_metadata(ap)
    display_title = _normalize_tag_value(meta.get("title")) or file_title
    metadata_text = _build_track_metadata_text(meta)
    s = RenderSettings(
        audio_path=ap,
        out_dir=base_settings.out_dir,
        resolution=base_settings.resolution,
        fps=base_settings.fps,
        style=base_settings.style,
        palette=base_settings.palette,
        background=base_settings.background,
        quality=base_settings.quality,
        glow=base_settings.glow,
        glow_intensity=base_settings.glow_intensity,
        blur=base_settings.blur,
        blur_amount=base_settings.blur_amount,
        particles=base_settings.particles,
        particle_count=base_settings.particle_count,
        mirror=base_settings.mirror,
        beat_react=base_settings.beat_react,
        logo_text=base_settings.logo_text,
        logo_img_path=base_settings.logo_img_path,
        logo_pos=base_settings.logo_pos,
        logo_scale=base_settings.logo_scale,
        logo_opacity=base_settings.logo_opacity,
        logo_start=base_settings.logo_start,
        logo_end=base_settings.logo_end,
        logo_follow_progress=base_settings.logo_follow_progress,
        bitrate=base_settings.bitrate,
        codec=base_settings.codec,
        draft_start=base_settings.draft_start,
        draft_duration=base_settings.draft_duration,
        progress_bar=base_settings.progress_bar,
        progress_pos=base_settings.progress_pos,
        progress_height=base_settings.progress_height,
        clock=base_settings.clock,
        clock_mode=base_settings.clock_mode,
        clock_style=base_settings.clock_style,
        clock_pos=base_settings.clock_pos,
        clock_scale=base_settings.clock_scale,
        clock_countdown=base_settings.clock_countdown,
        clock_offset_x=base_settings.clock_offset_x,
        clock_offset_y=base_settings.clock_offset_y,
        # Track title watermark settings
        track_title=display_title,
        track_title_watermark=base_settings.track_title_watermark,
        track_title_duration=base_settings.track_title_duration,
        track_title_pos=base_settings.track_title_pos,
        track_title_offset_x=base_settings.track_title_offset_x,
        track_title_offset_y=base_settings.track_title_offset_y,
        track_metadata_watermark=base_settings.track_metadata_watermark,
        track_metadata_duration=base_settings.track_metadata_duration,
        track_metadata_pos=base_settings.track_metadata_pos,
        track_metadata_offset_x=base_settings.track_metadata_offset_x,
        track_metadata_offset_y=base_settings.track_metadata_offset_y,
        track_metadata_text=metadata_text,
        # Background/Foreground image settings
        bg_image_path=base_settings.bg_image_path,
        bg_image_opacity=base_settings.bg_image_opacity,
        bg_image_pattern=base_settings.bg_image_pattern,
        fg_image_path=base_settings.fg_image_path,
        fg_image_opacity=base_settings.fg_image_opacity,
        # Font settings
        font_path=base_settings.font_path,
        font_size=base_settings.font_size,
        progress_cb=base_settings.progress_cb,
        stop_event=base_settings.stop_event,
    )
    out_path, dur = render_video(s)
    gc.collect()
    return out_path, file_title, dur


def render_batch(audio_paths: List[str], base_settings: RenderSettings, *, merge: bool = True, add_chapters: bool = True,
                 gap_duration: float = 0.0, gap_slate: bool = True, gap_label_format: str = "Next: {title}") -> List[str] | str:
    if not audio_paths:
        raise ValueError("No input files provided for batch")
    gap_duration = max(0.0, float(gap_duration))

    # cancellation hook
    if base_settings.stop_event is not None and base_settings.stop_event.is_set():
        raise RenderCancelled("Render cancelled by user")

    # Parallel per-file rendering
    workers = max(1, min(int(base_settings.workers), cpu_count(), len(audio_paths)))
    outs: List[str] = []
    titles: List[str] = []
    durations: List[float] = []
    slate_paths: List[str] = []
    tmp_files_to_cleanup: List[str] = []

    success = False
    try:
        if workers == 1:
            for i, ap in enumerate(audio_paths):
                if base_settings.stop_event is not None and base_settings.stop_event.is_set():
                    raise RenderCancelled("Render cancelled by user")
                if base_settings.progress_cb:
                    try:
                        base_settings.progress_cb(0, 1, f"batch:{i+1}/{len(audio_paths)}:{ap}")
                    except Exception:
                        pass
                out_path, title, dur = _render_one_for_batch(ap, base_settings)
                outs.append(out_path); titles.append(title); durations.append(dur)
        else:
            # We need to map futures back to original index to maintain order for merging/chapters.
            results: Dict[str, Tuple[str, str, float]] = {} # Map input path to result (path, title, dur)
            
            if base_settings.progress_cb:
                try:
                    base_settings.progress_cb(0, 1, f"batch:parallel:{len(audio_paths)}")
                except Exception:
                    pass
            with ThreadPoolExecutor(max_workers=workers) as ex:
                futs = {ex.submit(_render_one_for_batch, ap, base_settings): ap for ap in audio_paths}
                for fut in as_completed(futs):
                    original_path = futs[fut]
                    out_path, title, dur = fut.result()
                    results[original_path] = (out_path, title, dur)

            # Reconstruct ordered lists based on original audio_paths order
            for ap in audio_paths:
                if ap in results:
                    out_path, title, dur = results[ap]
                    outs.append(out_path)
                    titles.append(title)
                    durations.append(dur)
                else:
                    raise RuntimeError(f"Missing result for audio path: {ap}")

        if not merge:
            success = True
            return outs

        # Build concat list with optional slates
        list_file = os.path.join(_SESSION_TMPDIR, f"concat_list_{uuid.uuid4().hex}.txt")
        tmp_files_to_cleanup.append(list_file)
        width, height = base_settings.resolution
        fps = max(1, int(base_settings.fps))

        # Keep original order of inputs for concat
        with open(list_file, 'w', encoding='utf-8') as f:
            for i, out_path in enumerate(outs):
                f.write(f"file '{_escape_ffconcat_path(out_path)}'\n")
                if gap_duration > 0 and i < len(audio_paths) - 1:
                    next_title = titles[i+1]
                    if gap_slate:
                        label = (gap_label_format or "Next: {title}").format(index=i+2, title=next_title)
                        slate = _generate_slate_video(width, height, fps, gap_duration, label,
                                                      background=base_settings.background,
                                                      font_path=base_settings.font_path,
                                                      font_size=int(base_settings.font_size * 1.5))
                        slate_paths.append(slate)
                        f.write(f"file '{_escape_ffconcat_path(slate)}'\n")
                    else:
                        # If no slate label is requested, we still need a video segment for the gap duration
                        slate = _generate_slate_video(width, height, fps, gap_duration, "",
                                                      background=base_settings.background,
                                                      font_path=base_settings.font_path,
                                                      font_size=int(base_settings.font_size * 1.5))
                        slate_paths.append(slate)
                        f.write(f"file '{_escape_ffconcat_path(slate)}'\n")

        merged = os.path.join(base_settings.out_dir or os.path.dirname(outs[0]), "merged_with_chapters_temp.mp4")
        try:
            cmd_concat = ['ffmpeg', '-y', '-f', 'concat', '-safe', '0', '-i', list_file, '-c', 'copy', merged]
            subprocess.run(cmd_concat, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        except subprocess.CalledProcessError as e:
            err = (e.stderr or str(e)).strip()
            print(f"Error during ffmpeg concat: {err}")
            raise RuntimeError("FFmpeg concat failed. See log for details.") from e

        # Chapters and Tracklist generation
        if add_chapters:
            # Generate FFMETADATA (correctly calculating gaps)
            ffmeta_text = _ffmetadata_for_chapters(titles, durations, gap_duration)
            ffmeta_path = os.path.join(_SESSION_TMPDIR, f"chapters_{uuid.uuid4().hex}.ffmeta")
            tmp_files_to_cleanup.append(ffmeta_path)
            with open(ffmeta_path, 'w', encoding='utf-8') as f:
                f.write(ffmeta_text)
            
            # Add Tracklist file generation
            tracklist_path = os.path.join(base_settings.out_dir or os.path.dirname(outs[0]), 'tracklist.txt')
            _generate_tracklist_txt(titles, durations, gap_duration, tracklist_path)
            print(f"Tracklist saved: {tracklist_path}")
            
            # Mux chapters metadata into the final file
            final_path = os.path.join(base_settings.out_dir or os.path.dirname(outs[0]), 'merged_with_chapters.mp4')
            cmd_meta = ['ffmpeg', '-y', '-i', merged, '-i', ffmeta_path, '-map_metadata', '1', '-c', 'copy', final_path]
            try:
                subprocess.run(cmd_meta, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            except subprocess.CalledProcessError as e:
                err = (e.stderr or str(e)).strip()
                print(f"Error during metadata mux: {err}")
                raise RuntimeError("FFmpeg metadata mux failed.") from e
                
            try:
                os.remove(merged)
            except Exception:
                pass
            
            success = True
            return final_path

        success = True
        return merged

    finally:
        for _p in tmp_files_to_cleanup:
            try: os.remove(_p)
            except Exception: pass
            
        if (not success) or base_settings.cleanup_intermediates:
            for p in outs + slate_paths:
                try: os.remove(p)
                except Exception: pass