auto-editor 25.3.0__py3-none-any.whl → 26.0.0__py3-none-any.whl

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 = ""
@@ -112,35 +116,92 @@ class SubtitleParser:
         self.contents = new_content
 
     def write(self, file_path: str) -> None:
+        codec = self.codec
         with open(file_path, "w", encoding="utf-8") as file:
             file.write(self.header)
             for c in self.contents:
                 file.write(
-                    f"{c.before}{to_timecode(c.start / self.tb, self.codec)}"
-                    f"{c.middle}{to_timecode(c.end / self.tb, self.codec)}"
-                    f"{c.after}"
+                    f"{c.before}{to_timecode(c.start / self.tb, codec)}"
+                    + f"{c.middle}{to_timecode(c.end / self.tb, codec)}"
+                    + c.after
+                    + ("\n" if codec == "webvtt" else "")
                 )
             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, log: Log) -> 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 in ("webvtt", "ass", "ssa"):
+            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, log)
+        parser.parse(ret, sub.codec)
         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

auto_editor/render/video.py

@@ -1,26 +1,22 @@
 from __future__ import annotations
 
-import os.path
 from dataclasses import dataclass
-from subprocess import DEVNULL, PIPE
-from sys import platform
 from typing import TYPE_CHECKING
 
 import av
 import numpy as np
 
-from auto_editor.output import video_quality
+from auto_editor.output import parse_bitrate
 from auto_editor.timeline import TlImage, TlRect, TlVideo
-from auto_editor.utils.encoder import encoders
 from auto_editor.utils.types import color
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
+    from typing import Any
 
-    from auto_editor.ffwrapper import FFmpeg, FileInfo
+    from auto_editor.ffwrapper import FileInfo
     from auto_editor.timeline import v3
     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
 
@@ -90,8 +86,10 @@ def make_image_cache(tl: v3) -> dict[tuple[FileInfo, int], np.ndarray]:
 
 
 def render_av(
-    ffmpeg: FFmpeg, tl: v3, args: Args, bar: Bar, ctr: Container, log: Log
-) -> tuple[str, bool]:
+    output: av.container.OutputContainer, tl: v3, args: Args, bar: Bar, log: Log
+) -> Any:
+    from_ndarray = av.VideoFrame.from_ndarray
+
     src = tl.src
     cns: dict[FileInfo, av.container.InputContainer] = {}
     decoders: dict[FileInfo, Iterator[av.VideoFrame]] = {}
@@ -99,8 +97,8 @@ def render_av(
     tous: dict[FileInfo, int] = {}
 
     target_pix_fmt = "yuv420p"  # Reasonable default
+    target_fps = tl.tb  # Always constant
     img_cache = make_image_cache(tl)
-    temp = log.temp
 
     first_src: FileInfo | None = None
     for src in tl.sources:
@@ -133,14 +131,27 @@ def render_av(
     log.debug(f"Tous: {tous}")
     log.debug(f"Clips: {tl.v}")
 
-    target_pix_fmt = target_pix_fmt if target_pix_fmt in allowed_pix_fmt else "yuv420p"
-    log.debug(f"Target pix_fmt: {target_pix_fmt}")
-
-    apply_video_later = True
-    if args.video_codec in encoders:
-        apply_video_later = set(encoders[args.video_codec]).isdisjoint(allowed_pix_fmt)
+    if args.video_codec == "gif":
+        _c = av.Codec("gif", "w")
+        if _c.video_formats is not None and target_pix_fmt in (
+            f.name for f in _c.video_formats
+        ):
+            target_pix_fmt = target_pix_fmt
+        else:
+            target_pix_fmt = "rgb8"
+        del _c
+    else:
+        target_pix_fmt = (
+            target_pix_fmt if target_pix_fmt in allowed_pix_fmt else "yuv420p"
+        )
 
-    log.debug(f"apply video quality settings now: {not apply_video_later}")
+    ops = {"mov_flags": "faststart"}
+    output_stream = output.add_stream(args.video_codec, rate=target_fps, options=ops)
+    yield output_stream
+    if not isinstance(output_stream, av.VideoStream):
+        log.error(f"Not a known video codec: {args.video_codec}")
+    if src.videos and src.videos[0].lang is not None:
+        output_stream.metadata["language"] = src.videos[0].lang
 
     if args.scale == 1.0:
         target_width, target_height = tl.res
@@ -157,44 +168,33 @@ def render_av(
             scale_graph.add("buffersink"),
         )
 
-    spedup = os.path.join(temp, "spedup0.mp4")
-
-    cmd = [
-        "-hide_banner",
-        "-y",
-        "-f",
-        "rawvideo",
-        "-c:v",
-        "rawvideo",
-        "-pix_fmt",
-        target_pix_fmt,
-        "-s",
-        f"{target_width}*{target_height}",
-        "-framerate",
-        f"{tl.tb}",
-        "-i",
-        "-",
-        "-pix_fmt",
-        target_pix_fmt,
-    ]
-
-    if platform == "darwin":
-        # Fix videotoolbox issue with legacy macs
-        cmd += ["-allow_sw", "1"]
-
-    if apply_video_later:
-        cmd += ["-c:v", "mpeg4", "-qscale:v", "1"]
+    output_stream.width = target_width
+    output_stream.height = target_height
+    output_stream.pix_fmt = target_pix_fmt
+    output_stream.framerate = target_fps
+
+    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:
+        output_stream.color_range = color_range
+    if colorspace in (0, 1) or (colorspace >= 3 and colorspace < 16):
+        output_stream.colorspace = colorspace
+    if color_prim == 1 or (color_prim >= 4 and color_prim < 17):
+        output_stream.color_primaries = color_prim
+    if color_trc == 1 or (color_trc >= 4 and color_trc < 22):
+        output_stream.color_trc = color_trc
+
+    if args.video_bitrate != "auto":
+        output_stream.bit_rate = parse_bitrate(args.video_bitrate, log)
+        log.debug(f"video bitrate: {output_stream.bit_rate}")
     else:
-        cmd += video_quality(args)
+        log.debug(f"[auto] video bitrate: {output_stream.bit_rate}")
 
-    # Setting SAR requires re-encoding so we do it here.
     if src is not None and src.videos and (sar := src.videos[0].sar) is not None:
-        cmd.extend(["-vf", f"setsar={sar}"])
-
-    cmd.append(spedup)
-
-    process2 = ffmpeg.Popen(cmd, stdin=PIPE, stdout=DEVNULL, stderr=DEVNULL)
-    assert process2.stdin is not None
+        output_stream.sample_aspect_ratio = sar
 
     # First few frames can have an abnormal keyframe count, so never seek there.
     seek = 10
@@ -206,142 +206,128 @@ def render_av(
     bg = color(args.background)
     null_frame = make_solid(target_width, target_height, target_pix_fmt, bg)
     frame_index = -1
-    try:
-        for index in range(tl.end):
-            obj_list: list[VideoFrame | TlRect | TlImage] = []
-            for layer in tl.v:
-                for lobj in layer:
-                    if isinstance(lobj, TlVideo):
-                        if index >= lobj.start and index < (lobj.start + lobj.dur):
-                            _i = round((lobj.offset + index - lobj.start) * lobj.speed)
-                            obj_list.append(VideoFrame(_i, lobj.src))
-                    elif index >= lobj.start and index < lobj.start + lobj.dur:
-                        obj_list.append(lobj)
 
+    for index in range(tl.end):
+        obj_list: list[VideoFrame | TlRect | TlImage] = []
+        for layer in tl.v:
+            for lobj in layer:
+                if isinstance(lobj, TlVideo):
+                    if index >= lobj.start and index < (lobj.start + lobj.dur):
+                        _i = round((lobj.offset + index - lobj.start) * lobj.speed)
+                        obj_list.append(VideoFrame(_i, lobj.src))
+                elif index >= lobj.start and index < lobj.start + lobj.dur:
+                    obj_list.append(lobj)
+
+        if tl.v1 is not None:
+            # When there can be valid gaps in the timeline.
             frame = null_frame
-            for obj in obj_list:
-                if isinstance(obj, VideoFrame):
-                    my_stream = cns[obj.src].streams.video[0]
-                    if frame_index > obj.index:
-                        log.debug(f"Seek: {frame_index} -> 0")
-                        cns[obj.src].seek(0)
-                        try:
-                            frame = next(decoders[obj.src])
-                            frame_index = round(frame.time * tl.tb)
-                        except StopIteration:
-                            pass
-
-                    while frame_index < obj.index:
-                        # Check if skipping ahead is worth it.
-                        if (
-                            obj.index - frame_index > seek_cost[obj.src]
-                            and frame_index > seek
-                        ):
-                            seek = frame_index + (seek_cost[obj.src] // 2)
-                            seek_frame = frame_index
-                            log.debug(f"Seek: {frame_index} -> {obj.index}")
-                            cns[obj.src].seek(
-                                obj.index * tous[obj.src],
-                                stream=my_stream,
-                            )
-
-                        try:
-                            frame = next(decoders[obj.src])
-                            frame_index = round(frame.time * tl.tb)
-                        except StopIteration:
-                            log.debug(f"No source frame at {index=}. Using null frame")
-                            frame = null_frame
-                            break
-
-                        if seek_frame is not None:
-                            log.debug(
-                                f"Skipped {frame_index - seek_frame} frame indexes"
-                            )
-                            frames_saved += frame_index - seek_frame
-                            seek_frame = None
-                        if frame.key_frame:
-                            log.debug(f"Keyframe {frame_index} {frame.pts}")
-
-                    if (frame.width, frame.height) != tl.res:
-                        width, height = tl.res
-                        graph = av.filter.Graph()
-                        graph.link_nodes(
-                            graph.add_buffer(template=my_stream),
-                            graph.add(
-                                "scale",
-                                f"{width}:{height}:force_original_aspect_ratio=decrease:eval=frame",
-                            ),
-                            graph.add("pad", f"{width}:{height}:-1:-1:color={bg}"),
-                            graph.add("buffersink"),
-                        ).vpush(frame)
-                        frame = graph.vpull()
-                elif isinstance(obj, TlRect):
+        # else, use the last frame
+
+        for obj in obj_list:
+            if isinstance(obj, VideoFrame):
+                my_stream = cns[obj.src].streams.video[0]
+                if frame_index > obj.index:
+                    log.debug(f"Seek: {frame_index} -> 0")
+                    cns[obj.src].seek(0)
+                    try:
+                        frame = next(decoders[obj.src])
+                        frame_index = round(frame.time * tl.tb)
+                    except StopIteration:
+                        pass
+
+                while frame_index < obj.index:
+                    # Check if skipping ahead is worth it.
+                    if (
+                        obj.index - frame_index > seek_cost[obj.src]
+                        and frame_index > seek
+                    ):
+                        seek = frame_index + (seek_cost[obj.src] // 2)
+                        seek_frame = frame_index
+                        log.debug(f"Seek: {frame_index} -> {obj.index}")
+                        cns[obj.src].seek(obj.index * tous[obj.src], stream=my_stream)
+
+                    try:
+                        frame = next(decoders[obj.src])
+                        frame_index = round(frame.time * tl.tb)
+                    except StopIteration:
+                        log.debug(f"No source frame at {index=}. Using null frame")
+                        frame = null_frame
+                        break
+
+                    if seek_frame is not None:
+                        log.debug(f"Skipped {frame_index - seek_frame} frame indexes")
+                        frames_saved += frame_index - seek_frame
+                        seek_frame = None
+                    if frame.key_frame:
+                        log.debug(f"Keyframe {frame_index} {frame.pts}")
+
+                if (frame.width, frame.height) != tl.res:
+                    width, height = tl.res
                     graph = av.filter.Graph()
-                    x, y = obj.x, obj.y
                     graph.link_nodes(
                         graph.add_buffer(template=my_stream),
                         graph.add(
-                            "drawbox",
-                            f"x={x}:y={y}:w={obj.width}:h={obj.height}:color={obj.fill}:t=fill",
+                            "scale",
+                            f"{width}:{height}:force_original_aspect_ratio=decrease:eval=frame",
                         ),
+                        graph.add("pad", f"{width}:{height}:-1:-1:color={bg}"),
                         graph.add("buffersink"),
                     ).vpush(frame)
                     frame = graph.vpull()
-                elif isinstance(obj, TlImage):
-                    img = img_cache[(obj.src, obj.width)]
-                    array = frame.to_ndarray(format="rgb24")
-
-                    overlay_h, overlay_w, _ = img.shape
-                    x_pos, y_pos = obj.x, obj.y
-
-                    x_start = max(x_pos, 0)
-                    y_start = max(y_pos, 0)
-                    x_end = min(x_pos + overlay_w, frame.width)
-                    y_end = min(y_pos + overlay_h, frame.height)
-
-                    # Clip the overlay image to fit into the frame
-                    overlay_x_start = max(-x_pos, 0)
-                    overlay_y_start = max(-y_pos, 0)
-                    overlay_x_end = overlay_w - max(
-                        (x_pos + overlay_w) - frame.width, 0
-                    )
-                    overlay_y_end = overlay_h - max(
-                        (y_pos + overlay_h) - frame.height, 0
-                    )
-                    clipped_overlay = img[
-                        overlay_y_start:overlay_y_end, overlay_x_start:overlay_x_end
-                    ]
-
-                    # Create a region of interest (ROI) on the video frame
-                    roi = array[y_start:y_end, x_start:x_end]
-
-                    # Blend the overlay image with the ROI based on the opacity
-                    roi = (1 - obj.opacity) * roi + obj.opacity * clipped_overlay
-                    array[y_start:y_end, x_start:x_end] = roi
-                    array = np.clip(array, 0, 255).astype(np.uint8)
-
-                    frame = av.VideoFrame.from_ndarray(array, format="rgb24")
-
-            if scale_graph is not None and frame.width != target_width:
-                scale_graph.vpush(frame)
-                frame = scale_graph.vpull()
-
-            if frame.format.name != target_pix_fmt:
-                frame = frame.reformat(format=target_pix_fmt)
-                bar.tick(index)
-            elif index % 3 == 0:
-                bar.tick(index)
-
-            process2.stdin.write(frame.to_ndarray().tobytes())
-
-        bar.end()
-        process2.stdin.close()
-        process2.wait()
-    except (OSError, BrokenPipeError):
-        bar.end()
-        ffmpeg.run_check_errors(cmd, log, True)
-        log.error("FFmpeg Error!")
-
+            elif isinstance(obj, TlRect):
+                graph = av.filter.Graph()
+                x, y = obj.x, obj.y
+                graph.link_nodes(
+                    graph.add_buffer(template=my_stream),
+                    graph.add(
+                        "drawbox",
+                        f"x={x}:y={y}:w={obj.width}:h={obj.height}:color={obj.fill}:t=fill",
+                    ),
+                    graph.add("buffersink"),
+                ).vpush(frame)
+                frame = graph.vpull()
+            elif isinstance(obj, TlImage):
+                img = img_cache[(obj.src, obj.width)]
+                array = frame.to_ndarray(format="rgb24")
+
+                overlay_h, overlay_w, _ = img.shape
+                x_pos, y_pos = obj.x, obj.y
+
+                x_start = max(x_pos, 0)
+                y_start = max(y_pos, 0)
+                x_end = min(x_pos + overlay_w, frame.width)
+                y_end = min(y_pos + overlay_h, frame.height)
+
+                # Clip the overlay image to fit into the frame
+                overlay_x_start = max(-x_pos, 0)
+                overlay_y_start = max(-y_pos, 0)
+                overlay_x_end = overlay_w - max((x_pos + overlay_w) - frame.width, 0)
+                overlay_y_end = overlay_h - max((y_pos + overlay_h) - frame.height, 0)
+                clipped_overlay = img[
+                    overlay_y_start:overlay_y_end, overlay_x_start:overlay_x_end
+                ]
+
+                # Create a region of interest (ROI) on the video frame
+                roi = array[y_start:y_end, x_start:x_end]
+
+                # Blend the overlay image with the ROI based on the opacity
+                roi = (1 - obj.opacity) * roi + obj.opacity * clipped_overlay
+                array[y_start:y_end, x_start:x_end] = roi
+                array = np.clip(array, 0, 255).astype(np.uint8)
+
+                frame = from_ndarray(array, format="rgb24")
+
+        if scale_graph is not None and frame.width != target_width:
+            scale_graph.vpush(frame)
+            frame = scale_graph.vpull()
+
+        if frame.format.name != target_pix_fmt:
+            frame = frame.reformat(format=target_pix_fmt)
+            bar.tick(index)
+        elif index % 3 == 0:
+            bar.tick(index)
+
+        yield from_ndarray(frame.to_ndarray(), format=frame.format.name)
+
+    bar.end()
     log.debug(f"Total frames saved seeking: {frames_saved}")
-
-    return spedup, apply_video_later

auto_editor/subcommands/repl.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import sys
 from dataclasses import dataclass, field
 from fractions import Fraction
+from os import environ
 
 import auto_editor
 from auto_editor.analyze import Levels
@@ -71,11 +72,19 @@ def main(sys_args: list[str] = sys.argv[1:]) -> None:
     print(f"Auto-Editor {auto_editor.__version__}")
     text = None
 
+    no_color = bool(environ.get("NO_COLOR") or environ.get("AV_LOG_FORCE_NOCOLOR"))
+    if no_color:
+        bold_pink = bold_red = reset = ""
+    else:
+        bold_pink = "\033[1;95m"
+        bold_red = "\033[1;31m"
+        reset = "\033[0m"
+
     try:
         while True:
             try:
                 if text is None:
-                    text = input("> ")
+                    text = input(f"{bold_pink}>{reset} ")
                 else:
                     text += "\n" + input("   ")
             except KeyboardInterrupt as e:
@@ -97,8 +106,8 @@ def main(sys_args: list[str] = sys.argv[1:]) -> None:
 
             except ClosingError:
                 continue  # Allow user to continue adding text
-            except (MyError, ZeroDivisionError) as e:
-                print(f"error: {e}")
+            except MyError as e:
+                print(f"{bold_red}error{reset}: {e}")
 
             text = None