diff --git a/README.md b/README.md index fb328e9..d736e80 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,18 @@ with open("video.mp4", "rb") as f: pass ``` +### Slice into Clips + +```python +from simple_video_utils.slicing import slice_video + +# One MP4 (bytes) per (start, end) second range +clips = slice_video("video.mp4", [(0.0, 1.5), (2.0, 3.2)]) + +# Center-crop to a square and resize to 256x256 (e.g. for model input) +clips = slice_video("video.mp4", [(0.0, 1.5)], size=256) +``` + ### Remote Videos ```python diff --git a/simple_video_utils/slicing.py b/simple_video_utils/slicing.py new file mode 100644 index 0000000..241ce63 --- /dev/null +++ b/simple_video_utils/slicing.py @@ -0,0 +1,94 @@ +"""Cut a video into clips by time range. + +When no pixel change is needed — the target ``size`` is None, or the source is +already ``size`` x ``size`` and unrotated — packets are stream-copied: fast and +lossless, but the cut snaps to the keyframe at or before ``start``. Otherwise +frames are decoded, center-cropped to a square, resized, and re-encoded. +""" + +import io +from collections.abc import Iterator, Sequence +from fractions import Fraction + +import av +import numpy as np + +from simple_video_utils.frames import read_frames_exact +from simple_video_utils.metadata import video_metadata + + +def _center_crop_square(frame: np.ndarray) -> np.ndarray: + height, width = frame.shape[:2] + side = min(height, width) + top = (height - side) // 2 + left = (width - side) // 2 + return frame[top : top + side, left : left + side] + + +def _encode_clip(src: str, start: float, end: float, fps: float, size: int) -> bytes: + frames = list(read_frames_exact(src, start_time=start, end_time=end)) + if not frames: + return b"" + buffer = io.BytesIO() + with av.open(buffer, mode="w", format="mp4") as output: + stream = output.add_stream("h264", rate=Fraction(fps).limit_denominator(1000)) + stream.width = stream.height = size + stream.pix_fmt = "yuv420p" + for frame in frames: + video_frame = av.VideoFrame.from_ndarray(_center_crop_square(frame), format="rgb24") + output.mux(stream.encode(video_frame.reformat(width=size, height=size))) + output.mux(stream.encode()) + return buffer.getvalue() + + +def _copy_clip(src: str, start: float, end: float) -> bytes: + """Remux [start, end] seconds without re-encoding, cutting on the keyframe at/before start.""" + with av.open(src) as source: + in_stream = source.streams.video[0] + time_base = in_stream.time_base + origin = in_stream.start_time or 0 # pts is on the stream's absolute timeline + buffer = io.BytesIO() + first_dts = None + with av.open(buffer, mode="w", format="mp4") as output: + out_stream = output.add_stream_from_template(in_stream) + source.seek(int(start / time_base) + origin, stream=in_stream, backward=True) + end_pts = end / time_base + origin + for packet in source.demux(in_stream): + if packet.pts is None or packet.dts is None: + continue + if packet.pts > end_pts: + break + if first_dts is None: + first_dts = packet.dts + packet.pts -= first_dts + packet.dts -= first_dts + packet.stream = out_stream + output.mux(packet) + return buffer.getvalue() if first_dts is not None else b"" + + +def slice_video( + src: str, + slices: Sequence[tuple[float, float]], + size: int | None = None, +) -> Iterator[bytes]: + """Yield one MP4 clip (bytes) per (start, end) second range, in order. + + Yields lazily so a long slice list never holds every clip in memory. ``size`` + center-crops each frame to a square and resizes to ``size`` x ``size``; a + source that is already that size (and unrotated) is stream-copied instead. + Every slice must be within the video (``0 <= start <= end <= duration``); + an out-of-range or empty slice raises ``ValueError``. + """ + meta = video_metadata(src) + should_copy = size is None or (meta.width == meta.height == size and meta.rotation == 0) + + for start, end in slices: + if start < 0 or end <= start or (meta.duration is not None and end > meta.duration): + msg = f"slice ({start}, {end}) out of range [0, {meta.duration}]" + raise ValueError(msg) + clip = _copy_clip(src, start, end) if should_copy else _encode_clip(src, start, end, meta.fps, size) + if not clip: + msg = f"slice ({start}, {end}) has no frames" + raise ValueError(msg) + yield clip diff --git a/tests/test_slicing.py b/tests/test_slicing.py new file mode 100644 index 0000000..ccd4274 --- /dev/null +++ b/tests/test_slicing.py @@ -0,0 +1,72 @@ +import io + +import av +import numpy as np +import pytest + +from simple_video_utils.slicing import slice_video + + +def _write_video(path, width=320, height=240, frames=30, fps=30): + with av.open(str(path), mode="w") as container: + stream = container.add_stream("h264", rate=fps) + stream.width, stream.height, stream.pix_fmt = width, height, "yuv420p" + for i in range(frames): + arr = np.full((height, width, 3), i * 8 % 256, dtype=np.uint8) + for packet in stream.encode(av.VideoFrame.from_ndarray(arr, format="rgb24")): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + return str(path) + + +@pytest.fixture +def video(tmp_path): + return _write_video(tmp_path / "src.mp4") + + +def _dims(clip: bytes) -> tuple[int, int]: + with av.open(io.BytesIO(clip), mode="r") as container: + frame = next(container.decode(video=0)) + return frame.width, frame.height + + +def test_slice_returns_one_clip_per_range(video): + clips = list(slice_video(video, [(0.0, 0.3), (0.5, 0.8)])) + assert len(clips) == 2 + assert all(clip for clip in clips) + + +def test_slice_keeps_source_size_by_default(video): + # Default path stream-copies, so the source resolution is preserved. + [clip] = slice_video(video, [(0.0, 0.3)]) + assert _dims(clip) == (320, 240) + + +def test_slice_center_crops_and_resizes(video): + [clip] = slice_video(video, [(0.0, 0.3)], size=256) + assert _dims(clip) == (256, 256) + + +def test_matching_size_is_copied(tmp_path): + # Source is already 128x128, so size=128 needs no re-encode. + square = _write_video(tmp_path / "square.mp4", width=128, height=128, frames=15) + [clip] = slice_video(square, [(0.0, 0.3)], size=128) + assert _dims(clip) == (128, 128) + + +def test_out_of_range_slice_raises(video): + # Source is 1s; slices past the end, before 0, or reversed are errors. + for bad in [(5.0, 5.5), (-0.1, 0.3), (0.5, 0.2)]: + with pytest.raises(ValueError, match="out of range"): + list(slice_video(video, [bad])) + with pytest.raises(ValueError, match="out of range"): + list(slice_video(video, [bad], size=256)) + + +def test_zero_length_slice_raises(video): + # A clip needs positive duration; start == end is out of range in both paths. + with pytest.raises(ValueError, match="out of range"): + list(slice_video(video, [(0.5, 0.5)])) + with pytest.raises(ValueError, match="out of range"): + list(slice_video(video, [(0.5, 0.5)], size=256))