Skip to content

Commit

Permalink
Add encoders to bar title
Browse files Browse the repository at this point in the history
  • Loading branch information
WyattBlue committed Feb 19, 2025
1 parent 2b0bbec commit e66052b
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 13 deletions.
57 changes: 54 additions & 3 deletions auto_editor/edit.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,24 @@ def set_audio_codec(
return codec


def normalize_encoder(codec: str) -> tuple[str, bool]:
is_hardware = False
if "_" in codec:
is_hardware = codec.split("_")[-1] in {"videotoolbox", "cuvid", "nvec"}

id = av.Codec(codec, "w").id
name = {
27: "h264",
97: "gif",
173: "hevc",
167: "vp9",
139: "vp8",
86018: "aac",
}.get(id, codec)

return name, is_hardware


def parse_export(export: str, log: Log) -> dict[str, Any]:
exploded = export.split(":", maxsplit=1)
if len(exploded) == 1:
Expand Down Expand Up @@ -365,17 +383,46 @@ def make_media(tl: v3, output_path: str) -> None:

# Setup video
if ctr.default_vid != "none" and tl.v:
vframes = render_av(output, tl, args, bar, log)
vframes = render_av(output, tl, args, log)
output_stream = next(vframes)
else:
output_stream, vframes = None, iter([])

no_color = log.no_color or log.machine
encoder_titles = []
if output_stream is not None:
name, h = normalize_encoder(output_stream.name)
if no_color:
encoder_titles.append(name)
else:
is_bold = ";1" if h else ""
encoder_titles.append(f"\033[95{is_bold}m{name}")
if audio_streams:
name, h = normalize_encoder(audio_streams[0].name) # type: ignore
if no_color:
encoder_titles.append(name)
else:
is_bold = ";1" if h else ""
encoder_titles.append(f"\033[96{is_bold}m{name}")
if subtitle_streams:
name = subtitle_streams[0].name # type: ignore
if no_color:
encoder_titles.append(name)
else:
encoder_titles.append(f"\033[32m{name}")

title = f"({os.path.splitext(output_path)[1][1:]}) "
if no_color:
title += "+".join(encoder_titles)
else:
title += "\033[0m+".join(encoder_titles) + "\033[0m"
bar.start(tl.end, title)

# Process frames
while True:
audio_frames = [next(frames, None) for frames in audio_gen_frames]
video_frame = next(vframes, None)
index, video_frame = next(vframes, (0, None))
subtitle_frames = [next(packet, None) for packet in sub_gen_frames]

if (
all(frame is None for frame in audio_frames)
and video_frame is None
Expand Down Expand Up @@ -405,12 +452,16 @@ def make_media(tl: v3, output_path: str) -> None:
except av.FFmpegError as e:
log.error(e)

bar.tick(index)

# Flush streams
if output_stream is not None:
output.mux(output_stream.encode(None))
for audio_stream in audio_streams:
output.mux(audio_stream.encode(None))

bar.end()

# Close resources
for audio_input in audio_inputs:
audio_input.close()
Expand Down
13 changes: 3 additions & 10 deletions auto_editor/render/video.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@

from auto_editor.ffwrapper import FileInfo
from auto_editor.timeline import v3
from auto_editor.utils.bar import Bar
from auto_editor.utils.log import Log
from auto_editor.utils.types import Args

Expand All @@ -31,7 +30,7 @@ class VideoFrame:

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_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")
Expand Down Expand Up @@ -62,7 +61,7 @@ def make_image_cache(tl: v3) -> dict[tuple[FileInfo, int], np.ndarray]:


def render_av(
output: av.container.OutputContainer, tl: v3, args: Args, bar: Bar, log: Log
output: av.container.OutputContainer, tl: v3, args: Args, log: Log
) -> Any:
from_ndarray = av.VideoFrame.from_ndarray

Expand Down Expand Up @@ -191,8 +190,6 @@ def render_av(
seek_frame = None
frames_saved = 0

bar.start(tl.end, "Creating new video")

bg = args.background
null_frame = make_solid(target_width, target_height, target_pix_fmt, bg)
frame_index = -1
Expand Down Expand Up @@ -313,11 +310,7 @@ def render_av(

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)
yield (index, from_ndarray(frame.to_ndarray(), format=frame.format.name))

bar.end()
log.debug(f"Total frames saved seeking: {frames_saved}")

0 comments on commit e66052b

Please sign in to comment.