auto-editor 25.3.1__py3-none-any.whl → 26.0.1__py3-none-any.whl

auto_editor/output.py

@@ -2,26 +2,40 @@ from __future__ import annotations
 
 import os.path
 from dataclasses import dataclass, field
-from fractions import Fraction
 
 import av
 from av.audio.resampler import AudioResampler
 
-from auto_editor.ffwrapper import FFmpeg, FileInfo
+from auto_editor.ffwrapper import FileInfo
 from auto_editor.utils.bar import Bar
-from auto_editor.utils.container import Container
 from auto_editor.utils.log import Log
-from auto_editor.utils.types import Args
+from auto_editor.utils.types import _split_num_str
+
+
+def parse_bitrate(input_: str, log: Log) -> int:
+    try:
+        val, unit = _split_num_str(input_)
+    except Exception as e:
+        log.error(e)
+
+    if unit.lower() == "k":
+        return int(val * 1000)
+    if unit == "M":
+        return int(val * 1_000_000)
+    if unit == "G":
+        return int(val * 1_000_000_000)
+    if unit == "":
+        return int(val)
+
+    log.error(f"Unknown bitrate: {input_}")
 
 
 @dataclass(slots=True)
 class Ensure:
-    _ffmpeg: FFmpeg
     _bar: Bar
     _sr: int
     log: Log
     _audios: list[tuple[FileInfo, int]] = field(default_factory=list)
-    _subtitles: list[tuple[FileInfo, int, str]] = field(default_factory=list)
 
     def audio(self, src: FileInfo, stream: int) -> str:
         try:
@@ -52,193 +66,21 @@ class Ensure:
 
             bar.start(dur, "Extracting audio")
 
-            # PyAV always uses "stereo" layout, which is what we want.
-            output_astream = out_container.add_stream("pcm_s16le", rate=sample_rate)
-            assert isinstance(output_astream, av.audio.stream.AudioStream)
-
+            output_astream = out_container.add_stream(
+                "pcm_s16le", layout="stereo", rate=sample_rate
+            )
             resampler = AudioResampler(format="s16", layout="stereo", rate=sample_rate)
             for i, frame in enumerate(in_container.decode(astream)):
                 if i % 1500 == 0 and frame.time is not None:
                     bar.tick(frame.time)
 
                 for new_frame in resampler.resample(frame):
-                    for packet in output_astream.encode(new_frame):
-                        out_container.mux_one(packet)
+                    out_container.mux(output_astream.encode(new_frame))
 
-            for packet in output_astream.encode():
-                out_container.mux_one(packet)
+            out_container.mux(output_astream.encode(None))
 
             out_container.close()
             in_container.close()
             bar.end()
 
         return out_path
-
-    def subtitle(self, src: FileInfo, stream: int, ext: str) -> str:
-        try:
-            self._subtitles.index((src, stream, ext))
-            first_time = False
-        except ValueError:
-            self._subtitles.append((src, stream, ext))
-            first_time = True
-
-        out_path = os.path.join(self.log.temp, f"{stream}s.{ext}")
-
-        if first_time:
-            self.log.debug(f"Making external subtitle: {out_path}")
-            self.log.conwrite("Extracting subtitle")
-            self._ffmpeg.run(["-i", f"{src.path}", "-map", f"0:s:{stream}", out_path])
-
-        return out_path
-
-
-def _ffset(option: str, value: str | None) -> list[str]:
-    if value is None or value == "unset" or value == "reserved":
-        return []
-    return [option] + [value]
-
-
-def video_quality(args: Args) -> list[str]:
-    return (
-        _ffset("-b:v", args.video_bitrate)
-        + ["-c:v", args.video_codec]
-        + _ffset("-qscale:v", args.video_quality_scale)
-        + ["-movflags", "faststart"]
-    )
-
-
-def mux_quality_media(
-    ffmpeg: FFmpeg,
-    visual_output: list[tuple[bool, str]],
-    audio_output: list[str],
-    sub_output: list[str],
-    apply_v: bool,
-    ctr: Container,
-    output_path: str,
-    tb: Fraction,
-    args: Args,
-    src: FileInfo,
-    log: Log,
-) -> None:
-    v_tracks = len(visual_output)
-    a_tracks = len(audio_output)
-    s_tracks = 0 if args.sn else len(sub_output)
-
-    cmd = ["-hide_banner", "-y", "-i", f"{src.path}"]
-
-    same_container = src.path.suffix == os.path.splitext(output_path)[1]
-
-    for is_video, path in visual_output:
-        if is_video or ctr.allow_image:
-            cmd.extend(["-i", path])
-        else:
-            v_tracks -= 1
-
-    if a_tracks > 0:
-        if args.keep_tracks_separate and ctr.max_audios is None:
-            for path in audio_output:
-                cmd.extend(["-i", path])
-        else:
-            # Merge all the audio a_tracks into one.
-            new_a_file = os.path.join(log.temp, "new_audio.wav")
-            if a_tracks > 1:
-                new_cmd = []
-                for path in audio_output:
-                    new_cmd.extend(["-i", path])
-                new_cmd.extend(
-                    [
-                        "-filter_complex",
-                        f"amix=inputs={a_tracks}:duration=longest",
-                        "-ac",
-                        "2",
-                        new_a_file,
-                    ]
-                )
-                ffmpeg.run(new_cmd)
-                a_tracks = 1
-            else:
-                new_a_file = audio_output[0]
-            cmd.extend(["-i", new_a_file])
-
-    for subfile in sub_output:
-        cmd.extend(["-i", subfile])
-
-    for i in range(v_tracks + s_tracks + a_tracks):
-        cmd.extend(["-map", f"{i+1}:0"])
-
-    cmd.extend(["-map_metadata", "0"])
-
-    track = 0
-    for is_video, path in visual_output:
-        if is_video:
-            if apply_v:
-                cmd += video_quality(args)
-            else:
-                # Real video is only allowed on track 0
-                cmd += ["-c:v:0", "copy"]
-
-            if float(tb).is_integer():
-                cmd += ["-video_track_timescale", f"{tb}"]
-
-        elif ctr.allow_image:
-            ext = os.path.splitext(path)[1][1:]
-            cmd += [f"-c:v:{track}", ext, f"-disposition:v:{track}", "attached_pic"]
-
-        track += 1
-    del track
-
-    for i, vstream in enumerate(src.videos):
-        if i > v_tracks:
-            break
-        if vstream.lang is not None:
-            cmd.extend([f"-metadata:s:v:{i}", f"language={vstream.lang}"])
-    for i, astream in enumerate(src.audios):
-        if i > a_tracks:
-            break
-        if astream.lang is not None:
-            cmd.extend([f"-metadata:s:a:{i}", f"language={astream.lang}"])
-    for i, sstream in enumerate(src.subtitles):
-        if i > s_tracks:
-            break
-        if sstream.lang is not None:
-            cmd.extend([f"-metadata:s:s:{i}", f"language={sstream.lang}"])
-
-    if s_tracks > 0:
-        scodec = src.subtitles[0].codec
-        if same_container:
-            cmd.extend(["-c:s", scodec])
-        elif ctr.scodecs is not None:
-            if scodec not in ctr.scodecs:
-                scodec = ctr.default_sub
-            cmd.extend(["-c:s", scodec])
-
-    if a_tracks > 0:
-        cmd += _ffset("-c:a", args.audio_codec) + _ffset("-b:a", args.audio_bitrate)
-
-    if same_container and v_tracks > 0:
-        color_range = src.videos[0].color_range
-        colorspace = src.videos[0].color_space
-        color_prim = src.videos[0].color_primaries
-        color_trc = src.videos[0].color_transfer
-
-        if color_range == 1 or color_range == 2:
-            cmd.extend(["-color_range", f"{color_range}"])
-        if colorspace in (0, 1) or (colorspace >= 3 and colorspace < 16):
-            cmd.extend(["-colorspace", f"{colorspace}"])
-        if color_prim == 1 or (color_prim >= 4 and color_prim < 17):
-            cmd.extend(["-color_primaries", f"{color_prim}"])
-        if color_trc == 1 or (color_trc >= 4 and color_trc < 22):
-            cmd.extend(["-color_trc", f"{color_trc}"])
-
-    if args.extras is not None:
-        cmd.extend(args.extras.split(" "))
-    cmd.extend(["-strict", "-2"])  # Allow experimental codecs.
-
-    if s_tracks > 0:
-        cmd.extend(["-map", "0:t?"])  # Add input attachments to output.
-
-    if not args.dn:
-        cmd.extend(["-map", "0:d?"])
-
-    cmd.append(output_path)
-    ffmpeg.run_check_errors(cmd, path=output_path)

auto_editor/render/audio.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import io
 from pathlib import Path
 from platform import system
-from subprocess import PIPE
 
 import av
 import numpy as np
@@ -17,6 +16,7 @@ from auto_editor.output import Ensure
 from auto_editor.timeline import TlAudio, v3
 from auto_editor.utils.bar import Bar
 from auto_editor.utils.cmdkw import ParserError, parse_with_palet, pAttr, pAttrs
+from auto_editor.utils.container import Container
 from auto_editor.utils.log import Log
 from auto_editor.utils.types import Args
 from auto_editor.wavfile import AudioData, read, write
@@ -35,8 +35,6 @@ norm_types = {
     ),
 }
 
-file_null = "NUL" if system() in ("Windows", "cli") else "/dev/null"
-
 
 def parse_norm(norm: str, log: Log) -> dict | None:
     if norm == "#f":
@@ -58,7 +56,7 @@ def parse_norm(norm: str, log: Log) -> dict | None:
         log.error(e)
 
 
-def parse_ebu_bytes(norm: dict, stderr: bytes, log: Log) -> list[str]:
+def parse_ebu_bytes(norm: dict, stderr: bytes, log: Log) -> tuple[str, str]:
     start = end = 0
     lines = stderr.splitlines()
 
@@ -78,13 +76,7 @@ def parse_ebu_bytes(norm: dict, stderr: bytes, log: Log) -> list[str]:
     except MyError:
         log.error(f"Invalid loudnorm stats.\n{start=},{end=}\n{stderr!r}")
 
-    for key in (
-        "input_i",
-        "input_tp",
-        "input_lra",
-        "input_thresh",
-        "target_offset",
-    ):
+    for key in ("input_i", "input_tp", "input_lra", "input_thresh", "target_offset"):
         val = float(parsed[key])
         if val == float("-inf"):
             parsed[key] = -99
@@ -100,31 +92,12 @@ def parse_ebu_bytes(norm: dict, stderr: bytes, log: Log) -> list[str]:
     m_thresh = parsed["input_thresh"]
     target_offset = parsed["target_offset"]
 
-    return [
-        "-af",
-        f"loudnorm=i={norm['i']}:lra={norm['lra']}:tp={norm['tp']}:offset={target_offset}"
+    filter = (
+        f"i={norm['i']}:lra={norm['lra']}:tp={norm['tp']}:offset={target_offset}"
         f":measured_i={m_i}:measured_lra={m_lra}:measured_tp={m_tp}"
-        f":measured_thresh={m_thresh}:linear=true:print_format=json",
-    ]
-
-
-def parse_peak_bytes(t: float, stderr: bytes, log: Log) -> list[str]:
-    peak_level = None
-    for line in stderr.splitlines():
-        if line.startswith(b"[Parsed_astats_0") and b"Peak level dB:" in line:
-            try:
-                peak_level = float(line.split(b":")[1])
-            except Exception:
-                log.error(f"Invalid `astats` stats.\n{stderr!r}")
-            break
-
-    if peak_level is None:
-        log.error(f"Invalid `astats` stats.\n{stderr!r}")
-
-    adjustment = t - peak_level
-    log.debug(f"current peak level: {peak_level}")
-    log.print(f"peak adjustment: {adjustment}")
-    return ["-af", f"volume={adjustment}"]
+        f":measured_thresh={m_thresh}:linear=true:print_format=json"
+    )
+    return "loudnorm", filter
 
 
 def apply_audio_normalization(
@@ -135,13 +108,9 @@ def apply_audio_normalization(
             f"loudnorm=i={norm['i']}:lra={norm['lra']}:tp={norm['tp']}:"
             f"offset={norm['gain']}:print_format=json"
         )
-    else:
-        first_pass = "astats=measure_overall=Peak_level:measure_perchannel=0"
-
-    log.debug(f"audio norm first pass: {first_pass}")
-
-    stderr = ffmpeg.Popen(
-        [
+        log.debug(f"audio norm first pass: {first_pass}")
+        file_null = "NUL" if system() in ("Windows", "cli") else "/dev/null"
+        cmd = [
             "-hide_banner",
             "-i",
             f"{pre_master}",
@@ -152,19 +121,56 @@ def apply_audio_normalization(
             "-f",
             "null",
             file_null,
-        ],
-        stdin=PIPE,
-        stdout=PIPE,
-        stderr=PIPE,
-    ).communicate()[1]
-
-    if norm["tag"] == "ebu":
-        cmd = parse_ebu_bytes(norm, stderr, log)
+        ]
+        stderr = ffmpeg.Popen("EBU", cmd, log).communicate()[1]
+        name, filter_args = parse_ebu_bytes(norm, stderr, log)
     else:
         assert "t" in norm
-        cmd = parse_peak_bytes(norm["t"], stderr, log)
 
-    ffmpeg.run(["-i", f"{pre_master}"] + cmd + [f"{path}"])
+        def get_peak_level(frame: av.AudioFrame) -> float:
+            # Calculate peak level in dB
+            # Should be equivalent to: -af astats=measure_overall=Peak_level:measure_perchannel=0
+            max_amplitude = np.abs(frame.to_ndarray()).max()
+            if max_amplitude > 0.0:
+                return -20.0 * np.log10(max_amplitude)
+            return -99.0
+
+        with av.open(pre_master) as container:
+            max_peak_level = -99.0
+            assert len(container.streams.video) == 0
+            for frame in container.decode(audio=0):
+                peak_level = get_peak_level(frame)
+                max_peak_level = max(max_peak_level, peak_level)
+
+        adjustment = norm["t"] - max_peak_level
+        log.debug(f"current peak level: {max_peak_level}")
+        log.print(f"peak adjustment: {adjustment:.3f}dB")
+        name, filter_args = "volume", f"{adjustment}"
+
+    with av.open(pre_master) as container:
+        input_stream = container.streams.audio[0]
+
+        output_file = av.open(path, mode="w")
+        output_stream = output_file.add_stream("pcm_s16le", rate=input_stream.rate)
+
+        graph = av.filter.Graph()
+        graph.link_nodes(
+            graph.add_abuffer(template=input_stream),
+            graph.add(name, filter_args),
+            graph.add("abuffersink"),
+        ).configure()
+        for frame in container.decode(input_stream):
+            graph.push(frame)
+            while True:
+                try:
+                    aframe = graph.pull()
+                    assert isinstance(aframe, av.AudioFrame)
+                    output_file.mux(output_stream.encode(aframe))
+                except (av.BlockingIOError, av.EOFError):
+                    break
+
+        output_file.mux(output_stream.encode(None))
+        output_file.close()
 
 
 def process_audio_clip(
@@ -212,28 +218,109 @@ def process_audio_clip(
             try:
                 aframe = graph.pull()
                 assert isinstance(aframe, av.AudioFrame)
-                for packet in output_stream.encode(aframe):
-                    output_file.mux(packet)
+                output_file.mux(output_stream.encode(aframe))
             except (av.BlockingIOError, av.EOFError):
                 break
 
     # Flush the stream
-    for packet in output_stream.encode(None):
-        output_file.mux(packet)
+    output_file.mux(output_stream.encode(None))
 
     input_file.close()
     output_file.close()
 
     output_bytes.seek(0)
+    has_filesig = output_bytes.read(4)
+    output_bytes.seek(0)
+    if not has_filesig:  # Can rarely happen when clip is extremely small
+        return np.empty((0, 2), dtype=np.int16)
+
     return read(output_bytes)[1]
 
 
+def mix_audio_files(sr: int, audio_paths: list[str], output_path: str) -> None:
+    mixed_audio = None
+    max_length = 0
+
+    # First pass: determine the maximum length
+    for path in audio_paths:
+        container = av.open(path)
+        stream = container.streams.audio[0]
+
+        # Calculate duration in samples
+        assert stream.duration is not None
+        assert stream.time_base is not None
+        duration_samples = int(stream.duration * sr / stream.time_base.denominator)
+        max_length = max(max_length, duration_samples)
+        container.close()
+
+    # Second pass: read and mix audio
+    for path in audio_paths:
+        container = av.open(path)
+        stream = container.streams.audio[0]
+
+        resampler = av.audio.resampler.AudioResampler(
+            format="s16", layout="mono", rate=sr
+        )
+
+        audio_array: list[np.ndarray] = []
+        for frame in container.decode(audio=0):
+            frame.pts = None
+            resampled = resampler.resample(frame)[0]
+            audio_array.extend(resampled.to_ndarray().flatten())
+
+        # Pad or truncate to max_length
+        current_audio = np.array(audio_array[:max_length])
+        if len(current_audio) < max_length:
+            current_audio = np.pad(
+                current_audio, (0, max_length - len(current_audio)), "constant"
+            )
+
+        if mixed_audio is None:
+            mixed_audio = current_audio.astype(np.float32)
+        else:
+            mixed_audio += current_audio.astype(np.float32)
+
+        container.close()
+
+    if mixed_audio is None:
+        raise ValueError("mixed_audio is None")
+
+    # Normalize the mixed audio
+    max_val = np.max(np.abs(mixed_audio))
+    if max_val > 0:
+        mixed_audio = mixed_audio * (32767 / max_val)
+    mixed_audio = mixed_audio.astype(np.int16)  # type: ignore
+
+    output_container = av.open(output_path, mode="w")
+    output_stream = output_container.add_stream("pcm_s16le", rate=sr)
+
+    chunk_size = sr  # Process 1 second at a time
+    for i in range(0, len(mixed_audio), chunk_size):
+        # Shape becomes (1, samples) for mono
+        chunk = np.array([mixed_audio[i : i + chunk_size]])
+
+        frame = av.AudioFrame.from_ndarray(chunk, format="s16", layout="mono")
+        frame.rate = sr
+        frame.pts = i  # Set presentation timestamp
+
+        output_container.mux(output_stream.encode(frame))
+
+    output_container.mux(output_stream.encode(None))
+    output_container.close()
+
+
 def make_new_audio(
-    tl: v3, ensure: Ensure, args: Args, ffmpeg: FFmpeg, bar: Bar, log: Log
+    tl: v3,
+    ctr: Container,
+    ensure: Ensure,
+    args: Args,
+    ffmpeg: FFmpeg,
+    bar: Bar,
+    log: Log,
 ) -> list[str]:
     sr = tl.sr
     tb = tl.tb
-    output = []
+    output: list[str] = []
     samples: dict[tuple[FileInfo, int], AudioData] = {}
 
     norm = parse_norm(args.audio_normalize, log)
@@ -311,4 +398,9 @@ def make_new_audio(
         Path(temp, "asdf.map").unlink(missing_ok=True)
     except PermissionError:
         pass
+
+    if not (args.keep_tracks_separate and ctr.max_audios is None) and len(output) > 1:
+        new_a_file = f"{Path(temp, 'new_audio.wav')}"
+        mix_audio_files(sr, output, new_a_file)
+        return [new_a_file]
     return output

auto_editor/render/subtitle.py

@@ -1,18 +1,23 @@
 from __future__ import annotations
 
+import io
 import os
 import re
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
+import av
+
 from auto_editor.utils.func import to_timecode
 
 if TYPE_CHECKING:
     from fractions import Fraction
 
-    from auto_editor.output import Ensure
     from auto_editor.timeline import v3
     from auto_editor.utils.chunks import Chunks
+    from auto_editor.utils.log import Log
+
+    Input = av.container.InputContainer
 
 
 @dataclass(slots=True)
@@ -26,7 +31,6 @@ class SerialSub:
 
 class SubtitleParser:
     def __init__(self, tb: Fraction) -> None:
-        self.supported_codecs = ("ass", "webvtt", "mov_text")
         self.tb = tb
         self.contents: list[SerialSub] = []
         self.header = ""
@@ -125,24 +129,81 @@ class SubtitleParser:
             file.write(self.footer)
 
 
-def make_new_subtitles(tl: v3, ensure: Ensure, temp: str) -> list[str]:
+def make_srt(input_: Input, stream: int) -> str:
+    output_bytes = io.StringIO()
+    input_stream = input_.streams.subtitles[stream]
+    assert input_stream.time_base is not None
+    s = 1
+    for packet in input_.demux(input_stream):
+        if packet.dts is None or packet.pts is None or packet.duration is None:
+            continue
+
+        start = packet.pts * input_stream.time_base
+        end = start + packet.duration * input_stream.time_base
+
+        for subset in packet.decode():
+            start_time = to_timecode(start, "srt")
+            end_time = to_timecode(end, "srt")
+
+            sub = subset[0]
+            assert len(subset) == 1
+            assert isinstance(sub, av.subtitles.subtitle.AssSubtitle)
+
+            output_bytes.write(f"{s}\n{start_time} --> {end_time}\n")
+            output_bytes.write(sub.dialogue.decode("utf-8", errors="ignore") + "\n\n")
+            s += 1
+
+    output_bytes.seek(0)
+    return output_bytes.getvalue()
+
+
+def _ensure(input_: Input, format: str, stream: int) -> str:
+    output_bytes = io.BytesIO()
+    output = av.open(output_bytes, "w", format=format)
+
+    in_stream = input_.streams.subtitles[stream]
+    out_stream = output.add_stream(template=in_stream)
+
+    for packet in input_.demux(in_stream):
+        if packet.dts is None:
+            continue
+        packet.stream = out_stream
+        output.mux(packet)
+
+    output.close()
+    output_bytes.seek(0)
+    return output_bytes.getvalue().decode("utf-8", errors="ignore")
+
+
+def make_new_subtitles(tl: v3, log: Log) -> list[str]:
     if tl.v1 is None:
         return []
 
+    input_ = av.open(tl.v1.source.path)
     new_paths = []
 
     for s, sub in enumerate(tl.v1.source.subtitles):
-        new_path = os.path.join(temp, f"new{s}s.{sub.ext}")
-        parser = SubtitleParser(tl.tb)
+        if sub.codec == "mov_text":
+            continue
 
-        ext = sub.ext if sub.codec in parser.supported_codecs else "vtt"
-        file_path = ensure.subtitle(tl.v1.source, s, ext)
-
-        with open(file_path, encoding="utf-8") as file:
-            parser.parse(file.read(), sub.codec)
+        parser = SubtitleParser(tl.tb)
+        if sub.codec == "ssa":
+            format = "ass"
+        elif sub.codec in ("webvtt", "ass"):
+            format = sub.codec
+        else:
+            log.error(f"Unknown subtitle codec: {sub.codec}")
 
+        if sub.codec == "mov_text":
+            ret = make_srt(input_, s)
+        else:
+            ret = _ensure(input_, format, s)
+        parser.parse(ret, format)
         parser.edit(tl.v1.chunks)
+
+        new_path = os.path.join(log.temp, f"new{s}s.{sub.ext}")
         parser.write(new_path)
         new_paths.append(new_path)
 
+    input_.close()
     return new_paths