auto-editor 28.1.0__py3-none-any.whl → 29.0.1__py3-none-any.whl

auto_editor/render/subtitle.py

@@ -1,205 +0,0 @@
-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.timeline import v3
-    from auto_editor.utils.chunks import Chunks
-    from auto_editor.utils.log import Log
-
-    Input = av.container.InputContainer
-
-
-@dataclass(slots=True)
-class SerialSub:
-    start: int
-    end: int
-    before: str
-    middle: str
-    after: str
-
-
-class SubtitleParser:
-    def __init__(self, tb: Fraction) -> None:
-        self.tb = tb
-        self.contents: list[SerialSub] = []
-        self.header = ""
-        self.footer = ""
-
-    @staticmethod
-    def to_tick(text: str, codec: str, tb: Fraction) -> int:
-        boxes = text.replace(",", ".").split(":")
-        assert len(boxes) < 4
-
-        boxes.reverse()
-        multiply = (1, 60, 3600)
-        seconds = 0.0
-        for box, mul in zip(boxes, multiply):
-            seconds += float(box) * mul
-
-        return round(seconds * tb)
-
-    def parse(self, text: str, codec: str) -> None:
-        self.codec = codec
-        self.contents = []
-
-        if codec in {"ass", "ssa"}:
-            time_code = re.compile(r"(.*)(\d+:\d+:[\d.]+)(.*)(\d+:\d+:[\d.]+)(.*)")
-        elif codec == "webvtt":
-            time_code = re.compile(r"()(\d+:[\d.]+)( --> )(\d+:[\d.]+)(\n.*)")
-        elif codec == "mov_text":
-            time_code = re.compile(r"()(\d+:\d+:[\d,]+)( --> )(\d+:\d+:[\d,]+)(\n.*)")
-        else:
-            raise ValueError(f"codec {codec} not supported.")
-
-        i = 0
-        for reg in re.finditer(time_code, text):
-            i += 1
-            if i == 1:
-                self.header = text[: reg.span()[0]]
-
-            self.contents.append(
-                SerialSub(
-                    self.to_tick(reg.group(2), self.codec, self.tb),
-                    self.to_tick(reg.group(4), self.codec, self.tb),
-                    reg.group(1),
-                    reg.group(3),
-                    f"{reg.group(5)}\n",
-                )
-            )
-
-        if i == 0:
-            self.header = ""
-            self.footer = ""
-        else:
-            self.footer = text[reg.span()[1] :]
-
-    def edit(self, chunks: Chunks) -> None:
-        for cut in reversed(chunks):
-            the_speed = cut[2]
-            speed_factor = (
-                1 if (the_speed == 0 or the_speed >= 99999) else 1 - (1 / the_speed)
-            )
-
-            new_content = []
-            for content in self.contents:
-                if cut[0] <= content.end and cut[1] > content.start:
-                    diff = int(
-                        (min(cut[1], content.end) - max(cut[0], content.start))
-                        * speed_factor
-                    )
-                    if content.start > cut[0]:
-                        content.start -= diff
-                        content.end -= diff
-
-                    content.end -= diff
-
-                elif content.start >= cut[0]:
-                    diff = int((cut[1] - cut[0]) * speed_factor)
-
-                    content.start -= diff
-                    content.end -= diff
-
-                if content.start != content.end:
-                    new_content.append(content)
-
-        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, codec)}"
-                    + f"{c.middle}{to_timecode(c.end / self.tb, codec)}"
-                    + c.after
-                    + ("\n" if codec == "webvtt" else "")
-                )
-            file.write(self.footer)
-
-
-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_num = packet.pts * input_stream.time_base
-        start = to_timecode(start_num, "srt")
-        end = to_timecode(start_num + packet.duration * input_stream.time_base, "srt")
-
-        for sub in packet.decode():
-            assert isinstance(sub, av.subtitles.subtitle.AssSubtitle)
-
-            output_bytes.write(f"{s}\n{start} --> {end}\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_from_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):
-        if sub.codec == "mov_text":
-            continue
-
-        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

auto_editor/render/video.py

@@ -1,307 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import TYPE_CHECKING
-
-import av
-import numpy as np
-
-from auto_editor.timeline import Clip, TlImage, TlRect
-from auto_editor.utils.func import parse_bitrate
-
-if TYPE_CHECKING:
-    from collections.abc import Iterator
-
-    from auto_editor.__main__ import Args
-    from auto_editor.ffwrapper import FileInfo
-    from auto_editor.timeline import v3
-    from auto_editor.utils.log import Log
-
-
-@dataclass(slots=True)
-class VideoFrame:
-    index: int
-    src: FileInfo
-
-
-def make_solid(width: int, height: int, pix_fmt: str, bg: str) -> av.VideoFrame:
-    hex_color = bg.lstrip("#").upper()
-    rgb_color = tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4))
-
-    rgb_array = np.full((height, width, 3), rgb_color, dtype=np.uint8)
-    rgb_frame = av.VideoFrame.from_ndarray(rgb_array, format="rgb24")
-    return rgb_frame.reformat(format=pix_fmt)
-
-
-def make_image_cache(tl: v3) -> dict[tuple[FileInfo, int], np.ndarray]:
-    img_cache = {}
-    for clip in tl.v:
-        for obj in clip:
-            if isinstance(obj, TlImage) and (obj.src, obj.width) not in img_cache:
-                with av.open(obj.src.path) as cn:
-                    my_stream = cn.streams.video[0]
-                    for frame in cn.decode(my_stream):
-                        if obj.width != 0:
-                            graph = av.filter.Graph()
-                            graph.link_nodes(
-                                graph.add_buffer(template=my_stream),
-                                graph.add("scale", f"{obj.width}:-1"),
-                                graph.add("buffersink"),
-                            ).vpush(frame)
-                            frame = graph.vpull()
-                        img_cache[(obj.src, obj.width)] = frame.to_ndarray(
-                            format="rgb24"
-                        )
-                        break
-    return img_cache
-
-
-def render_av(
-    output: av.container.OutputContainer, tl: v3, args: Args, log: Log
-) -> Iterator[tuple[int, av.VideoFrame]]:
-    from_ndarray = av.VideoFrame.from_ndarray
-
-    cns: dict[FileInfo, av.container.InputContainer] = {}
-    decoders: dict[FileInfo, Iterator[av.VideoFrame]] = {}
-    seek_cost: dict[FileInfo, int] = {}
-    tous: dict[FileInfo, int] = {}
-
-    pix_fmt = "yuv420p"  # Reasonable default
-    target_fps = tl.tb  # Always constant
-    img_cache = make_image_cache(tl)
-
-    first_src: FileInfo | None = None
-    for src in tl.sources:
-        if first_src is None:
-            first_src = src
-
-        if src not in cns:
-            cns[src] = av.open(f"{src.path}")
-
-    for src, cn in cns.items():
-        if len(cn.streams.video) > 0:
-            stream = cn.streams.video[0]
-            stream.thread_type = "AUTO"
-
-            if args.no_seek or stream.average_rate is None or stream.time_base is None:
-                sc_val = 4294967295  # 2 ** 32 - 1
-                tou = 0
-            else:
-                # Keyframes are usually spread out every 5 seconds or less.
-                sc_val = int(stream.average_rate * 5)
-                tou = int(stream.time_base.denominator / stream.average_rate)
-
-            seek_cost[src] = sc_val
-            tous[src] = tou
-            decoders[src] = cn.decode(stream)
-
-            if src == first_src and stream.pix_fmt is not None:
-                pix_fmt = stream.pix_fmt
-
-    log.debug(f"Tous: {tous}")
-    log.debug(f"Clips: {tl.v}")
-
-    output_stream = output.add_stream(args.video_codec, rate=target_fps)
-    output_stream.options = {"x265-params": "log-level=error"}
-
-    need_valid_fmt = not (
-        output_stream.codec.video_formats
-        and any(pix_fmt == vfmt.name for vfmt in output_stream.codec.video_formats)
-    )
-    if need_valid_fmt:
-        match output_stream.codec.canonical_name:
-            case "gif":
-                pix_fmt = "rgb8"
-            case "prores":
-                pix_fmt = "yuv422p10le"
-            case _:
-                pix_fmt = "yuv420p"
-
-    cc = output_stream.codec_context
-    if args.vprofile is not None:
-        if args.vprofile.title() not in cc.profiles:
-            b = " ".join([f'"{x.lower()}"' for x in cc.profiles])
-            log.error(
-                f"`{args.vprofile}` is not a valid profile.\nprofiles supported: {b}"
-            )
-
-        cc.profile = args.vprofile.title()
-
-    yield output_stream  # type: ignore
-    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
-        scale_graph = None
-    else:
-        target_width = max(round(tl.res[0] * args.scale), 2)
-        target_height = max(round(tl.res[1] * args.scale), 2)
-        scale_graph = av.filter.Graph()
-        scale_graph.link_nodes(
-            scale_graph.add(
-                "buffer", video_size="1x1", time_base="1/1", pix_fmt=pix_fmt
-            ),
-            scale_graph.add("scale", f"{target_width}:{target_height}"),
-            scale_graph.add("buffersink"),
-        )
-
-    output_stream.width = target_width
-    output_stream.height = target_height
-    output_stream.pix_fmt = 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 in {1, 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:
-        log.debug(f"[auto] video bitrate: {output_stream.bit_rate}")
-
-    if src is not None and src.videos and (sar := src.videos[0].sar) is not None:
-        output_stream.sample_aspect_ratio = sar
-
-    # First few frames can have an abnormal keyframe count, so never seek there.
-    seek = 10
-    seek_frame = None
-    frames_saved = 0
-
-    bg = args.background
-    null_frame = make_solid(target_width, target_height, pix_fmt, bg)
-    frame_index = -1
-
-    for index in range(tl.end):
-        obj_list: list[VideoFrame | TlRect | TlImage] = []
-        for layer in tl.v:
-            for lobj in layer:
-                if isinstance(lobj, Clip):
-                    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
-        # 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()
-                    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):
-                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  # type: ignore
-                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()
-
-        frame = frame.reformat(format=pix_fmt)
-        frame.pts = None  # type: ignore
-        frame.time_base = 0  # type: ignore
-        yield (index, frame)
-
-    log.debug(f"Total frames saved seeking: {frames_saved}")