################################################################################
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################
"""Video codec helpers backed by Java-side Flink FileSystem reads."""
from typing import Optional
from pyflink.common import Row
from pyflink.common.image import Image, ImageMode
from pyflink.fn_execution.file_system_bridge import (
FlinkFileSystemBridgeClient,
JavaBackedSeekableFile,
)
from pyflink.multimodal.types.video import VideoRef
_VALID_FRAME_SELECTORS = {"keyframe", "sample", "all_frames"}
_DEFAULT_SAMPLE_INTERVAL_MS = 1000
[docs]def open_video_container(uri: str,
bridge_client: FlinkFileSystemBridgeClient,
*,
container_options: Optional[dict] = None,
read_chunk_size: Optional[int] = None,
max_cached_blocks: Optional[int] = None,
read_ahead_blocks: Optional[int] = None):
"""
Opens a video container with PyAV using a Java-backed seekable file.
The returned object delegates to the PyAV container. The caller owns
closing it; close releases both the PyAV container and the Java-side file
handle.
"""
try:
import av
except ImportError as e:
raise ImportError(
"PyAV is required for video decoding. Install it with `pip install av`.") from e
file_obj = JavaBackedSeekableFile(
bridge_client,
uri,
**_file_read_options(read_chunk_size, max_cached_blocks, read_ahead_blocks))
try:
# Keep PyAV hardware acceleration out of part1. This pipeline
# materializes decoded frames as CPU ndarrays/IMAGE values, so
# hwaccel would still copy frames back to host memory and needs a
# separate benchmark plus Flink GPU resource scheduling design.
return _JavaBackedVideoContainer(
av.open(file_obj, **(container_options or {})), file_obj)
except Exception:
_close_resource(file_obj)
raise
class _JavaBackedVideoContainer(object):
"""PyAV container wrapper that owns its Java-backed file object."""
def __init__(self, container, file_obj):
self._container = container
self._file_obj = file_obj
def __getattr__(self, name):
return getattr(self._container, name)
def close(self):
try:
_close_resource(self._container)
finally:
_close_resource(self._file_obj)
def _normalize_video_ref(video_input):
if video_input is None:
return None
if isinstance(video_input, VideoRef):
return video_input
if isinstance(video_input, str):
return VideoRef(video_input)
uri = getattr(video_input, "uri", None)
if uri is not None:
return VideoRef(
uri=uri,
start_time_ms=_optional_time_ms(
"start_time_ms", getattr(video_input, "start_time_ms", None)),
end_time_ms=_optional_time_ms(
"end_time_ms", getattr(video_input, "end_time_ms", None)),
)
try:
uri, start_time_ms, end_time_ms = video_input[0], video_input[1], video_input[2]
except (TypeError, IndexError) as e:
raise ValueError(
"video_input must be a URI string, VideoRef, or VIDEO_STRUCT row"
) from e
return VideoRef(
uri=uri,
start_time_ms=_optional_time_ms("start_time_ms", start_time_ms),
end_time_ms=_optional_time_ms("end_time_ms", end_time_ms),
)
def _optional_int(value):
return None if value is None else int(value)
def _optional_time_ms(name, value):
if value is None:
return None
if not isinstance(value, int) or isinstance(value, bool) or value < 0:
raise ValueError(f"{name} must be a non-negative integer, got {value!r}")
return value
def _optional_float(value):
return None if value is None else float(value)
def _optional_str(value):
return None if value is None else str(value)
def _normalize_frame_selector(frame_selector):
if frame_selector not in _VALID_FRAME_SELECTORS:
raise ValueError(
f"frame_selector must be one of {sorted(_VALID_FRAME_SELECTORS)}, "
f"got {frame_selector!r}"
)
return frame_selector
def _validate_positive_int(name, value):
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
raise ValueError(f"{name} must be a positive integer, got {value!r}")
return value
def _validate_optional_positive_int(name, value):
if value is None:
return None
return _validate_positive_int(name, value)
def _file_read_options(read_chunk_size, max_cached_blocks, read_ahead_blocks):
options = {}
if read_chunk_size is not None:
options["read_chunk_size"] = read_chunk_size
if max_cached_blocks is not None:
options["max_cached_blocks"] = max_cached_blocks
if read_ahead_blocks is not None:
options["read_ahead_blocks"] = read_ahead_blocks
return options
def _close_resource(resource):
close = getattr(resource, "close", None)
if close is not None:
close()
def _first_video_stream(container):
streams = getattr(container, "streams", None)
video_streams = getattr(streams, "video", None) if streams is not None else None
if not video_streams:
raise ValueError("Video container has no video stream.")
return video_streams[0]
def _video_duration_ms(container, stream):
stream_duration_ms = _duration_from_stream(stream)
if stream_duration_ms is not None:
return stream_duration_ms
container_duration = getattr(container, "duration", None)
if container_duration is None:
return None
return int(round(int(container_duration) / 1000.0))
def _duration_from_stream(stream):
duration = getattr(stream, "duration", None)
time_base = getattr(stream, "time_base", None)
if duration is None or not time_base:
return None
return int(round(int(duration) * float(time_base) * 1000))
def _stream_fps(stream):
average_rate = getattr(stream, "average_rate", None)
if average_rate is not None:
return _optional_float(average_rate)
base_rate = getattr(stream, "base_rate", None)
return _optional_float(base_rate)
def _stream_codec_name(stream):
codec_context = getattr(stream, "codec_context", None)
codec_name = getattr(codec_context, "name", None)
if codec_name is not None:
return _optional_str(codec_name)
codec = getattr(stream, "codec", None)
return _optional_str(getattr(codec, "name", None))
def _seek_to_segment_start(container, stream, start_time_ms):
if start_time_ms is None or start_time_ms <= 0:
return
offset = _stream_offset_from_ms(stream, start_time_ms)
try:
container.seek(offset, any_frame=False, backward=True, stream=stream)
except (AttributeError, OSError, ValueError):
container.seek(start_time_ms * 1000, any_frame=False, backward=True)
def _stream_offset_from_ms(stream, time_ms):
time_base = getattr(stream, "time_base", None)
if time_base:
return int(round((time_ms / 1000.0) / float(time_base)))
return time_ms * 1000
def _frame_time_ms(frame, stream):
frame_time = getattr(frame, "time", None)
if frame_time is not None:
return int(round(float(frame_time) * 1000))
pts = getattr(frame, "pts", None)
time_base = getattr(stream, "time_base", None)
if pts is not None and time_base:
return int(round(int(pts) * float(time_base) * 1000))
return None
def _before_segment(time_ms, start_time_ms):
return time_ms is not None and start_time_ms is not None and time_ms < start_time_ms
def _after_segment(time_ms, end_time_ms):
return time_ms is not None and end_time_ms is not None and time_ms > end_time_ms
def _frame_to_output(video_ref, frame, stream, frame_index, time_ms):
pixel_array = frame.to_ndarray(format="rgb24")
image = Image(data=pixel_array, mode=ImageMode.RGB)
metadata = Row(
uri=video_ref.uri,
video_stream_index=_optional_int(getattr(stream, "index", None)),
frame_index=frame_index,
pts=_optional_int(getattr(frame, "pts", None)),
time_ms=time_ms,
key_frame=bool(getattr(frame, "key_frame", False)),
start_time_ms=video_ref.start_time_ms,
end_time_ms=video_ref.end_time_ms,
)
return image, metadata