105 lines
3.8 KiB
Python
105 lines
3.8 KiB
Python
"""프레임 PNG를 디스크에 풀지 않고 메모리에서 디코드·인코드하는 비디오 유틸."""
|
|
import io
|
|
from collections.abc import Iterable, Iterator, Sequence
|
|
from fractions import Fraction
|
|
from pathlib import Path
|
|
|
|
import av
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
AAC_FRAME_SIZE = 1024
|
|
|
|
|
|
def open_clip(clip: bytes | Path):
|
|
return av.open(io.BytesIO(clip) if isinstance(clip, bytes) else str(clip))
|
|
|
|
|
|
def duration(clip: bytes | Path) -> float:
|
|
container = open_clip(clip)
|
|
try:
|
|
if container.duration is None:
|
|
raise RuntimeError("클립 길이를 읽을 수 없다")
|
|
return container.duration / av.time_base
|
|
finally:
|
|
container.close()
|
|
|
|
|
|
def frames_at(clip: bytes | Path, times: Sequence[float]) -> list[Image.Image]:
|
|
"""지정한 시각들의 프레임. seek은 키프레임 단위라 그 이후 첫 프레임을 취한다."""
|
|
container = open_clip(clip)
|
|
stream = container.streams.video[0]
|
|
frames = []
|
|
try:
|
|
for at in times:
|
|
container.seek(int(at / stream.time_base), stream=stream)
|
|
picked = None
|
|
for frame in container.decode(stream):
|
|
picked = frame
|
|
if frame.time is not None and frame.time >= at - 1e-3:
|
|
break
|
|
if picked is None:
|
|
raise RuntimeError(f"t={at}s 프레임을 못 뽑았다")
|
|
frames.append(picked.to_image())
|
|
finally:
|
|
container.close()
|
|
return frames
|
|
|
|
|
|
def decode_frames(clip: bytes | Path) -> Iterator[Image.Image]:
|
|
"""모든 프레임을 순서대로 흘려보낸다. 한 장씩만 메모리에 올라간다."""
|
|
container = open_clip(clip)
|
|
try:
|
|
for frame in container.decode(video=0):
|
|
yield frame.to_image()
|
|
finally:
|
|
container.close()
|
|
|
|
|
|
def encode_mp4(frames: Iterable[Image.Image], size: tuple[int, int],
|
|
audio: np.ndarray | None = None, *, fps: int = 24,
|
|
sample_rate: int = 44100, crf: int = 18,
|
|
audio_bit_rate: int = 192000) -> bytes:
|
|
"""H.264 영상 + AAC 스테레오 오디오를 mp4 바이트로 뽑는다. 모노 오디오는 양 채널에 복제한다."""
|
|
buffer = io.BytesIO()
|
|
container = av.open(buffer, "w", format="mp4")
|
|
|
|
video_stream = container.add_stream("libx264", rate=fps)
|
|
video_stream.width, video_stream.height = size
|
|
video_stream.pix_fmt = "yuv420p"
|
|
video_stream.options = {"crf": str(crf), "preset": "medium"}
|
|
|
|
audio_stream = None
|
|
if audio is not None:
|
|
audio_stream = container.add_stream("aac", rate=sample_rate)
|
|
audio_stream.layout = "stereo"
|
|
audio_stream.bit_rate = audio_bit_rate
|
|
|
|
for index, image in enumerate(frames):
|
|
frame = av.VideoFrame.from_image(image)
|
|
frame.pts = index
|
|
for packet in video_stream.encode(frame):
|
|
container.mux(packet)
|
|
for packet in video_stream.encode():
|
|
container.mux(packet)
|
|
|
|
if audio_stream is not None:
|
|
interleaved = np.repeat(audio, 2)
|
|
block_size = audio_stream.frame_size or AAC_FRAME_SIZE
|
|
for start in range(0, len(audio), block_size):
|
|
block = interleaved[start * 2:(start + block_size) * 2]
|
|
if len(block) < block_size * 2:
|
|
block = np.pad(block, (0, block_size * 2 - len(block)))
|
|
frame = av.AudioFrame.from_ndarray(block.reshape(1, -1),
|
|
format="s16", layout="stereo")
|
|
frame.rate = sample_rate
|
|
frame.pts = start
|
|
frame.time_base = Fraction(1, sample_rate)
|
|
for packet in audio_stream.encode(frame):
|
|
container.mux(packet)
|
|
for packet in audio_stream.encode():
|
|
container.mux(packet)
|
|
|
|
container.close()
|
|
return buffer.getvalue()
|