File size: 11,694 Bytes
f1e6b80 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 |
import os
from dataclasses import dataclass, field
from io import BytesIO
from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union
import numpy as np
import pyarrow as pa
from .. import config
from ..download.download_config import DownloadConfig
from ..table import array_cast
from ..utils.file_utils import is_local_path, xopen
from ..utils.py_utils import string_to_dict
if TYPE_CHECKING:
from decord import VideoReader
from .features import FeatureType
@dataclass
class Video:
"""
**Experimental.** Video [`Feature`] to read video data from a video file.
Input: The Video feature accepts as input:
- A `str`: Absolute path to the video file (i.e. random access is allowed).
- A `dict` with the keys:
- `path`: String with relative path of the video file in a dataset repository.
- `bytes`: Bytes of the video file.
This is useful for archived files with sequential access.
- A `decord.VideoReader`: decord video reader object.
Args:
mode (`str`, *optional*):
The mode to convert the video to. If `None`, the native mode of the video is used.
decode (`bool`, defaults to `True`):
Whether to decode the video data. If `False`,
returns the underlying dictionary in the format `{"path": video_path, "bytes": video_bytes}`.
Examples:
```py
>>> from datasets import Dataset, Video
>>> ds = Dataset.from_dict({"video":["path/to/Screen Recording.mov"]}).cast_column("video", Video())
>>> ds.features["video"]
Video(decode=True, id=None)
>>> ds[0]["video"]
<decord.video_reader.VideoReader at 0x105525c70>
>>> ds = ds.cast_column('video', Video(decode=False))
{'bytes': None,
'path': 'path/to/Screen Recording.mov'}
```
"""
decode: bool = True
id: Optional[str] = None
# Automatically constructed
dtype: ClassVar[str] = "decord.VideoReader"
pa_type: ClassVar[Any] = pa.struct({"bytes": pa.binary(), "path": pa.string()})
_type: str = field(default="Video", init=False, repr=False)
def __post_init__(self):
if config.DECORD_AVAILABLE:
patch_decord()
def __call__(self):
return self.pa_type
def encode_example(self, value: Union[str, bytes, dict, np.ndarray, "VideoReader"]) -> dict:
"""Encode example into a format for Arrow.
Args:
value (`str`, `np.ndarray`, `VideoReader` or `dict`):
Data passed as input to Video feature.
Returns:
`dict` with "path" and "bytes" fields
"""
if config.DECORD_AVAILABLE:
from decord import VideoReader
else:
VideoReader = None
if isinstance(value, list):
value = np.array(value)
if isinstance(value, str):
return {"path": value, "bytes": None}
elif isinstance(value, bytes):
return {"path": None, "bytes": value}
elif isinstance(value, np.ndarray):
# convert the video array to bytes
return encode_np_array(value)
elif VideoReader and isinstance(value, VideoReader):
# convert the decord video reader to bytes
return encode_decord_video(value)
elif value.get("path") is not None and os.path.isfile(value["path"]):
# we set "bytes": None to not duplicate the data if they're already available locally
return {"bytes": None, "path": value.get("path")}
elif value.get("bytes") is not None or value.get("path") is not None:
# store the video bytes, and path is used to infer the video format using the file extension
return {"bytes": value.get("bytes"), "path": value.get("path")}
else:
raise ValueError(
f"A video sample should have one of 'path' or 'bytes' but they are missing or None in {value}."
)
def decode_example(self, value: dict, token_per_repo_id=None) -> "VideoReader":
"""Decode example video file into video data.
Args:
value (`str` or `dict`):
A string with the absolute video file path, a dictionary with
keys:
- `path`: String with absolute or relative video file path.
- `bytes`: The bytes of the video file.
token_per_repo_id (`dict`, *optional*):
To access and decode
video files from private repositories on the Hub, you can pass
a dictionary repo_id (`str`) -> token (`bool` or `str`).
Returns:
`decord.VideoReader`
"""
if not self.decode:
raise RuntimeError("Decoding is disabled for this feature. Please use Video(decode=True) instead.")
if config.DECORD_AVAILABLE:
from decord import VideoReader
else:
raise ImportError("To support decoding videos, please install 'decord'.")
if token_per_repo_id is None:
token_per_repo_id = {}
path, bytes_ = value["path"], value["bytes"]
if bytes_ is None:
if path is None:
raise ValueError(f"A video should have one of 'path' or 'bytes' but both are None in {value}.")
else:
if is_local_path(path):
video = VideoReader(path)
else:
source_url = path.split("::")[-1]
pattern = (
config.HUB_DATASETS_URL
if source_url.startswith(config.HF_ENDPOINT)
else config.HUB_DATASETS_HFFS_URL
)
try:
repo_id = string_to_dict(source_url, pattern)["repo_id"]
token = token_per_repo_id.get(repo_id)
except ValueError:
token = None
download_config = DownloadConfig(token=token)
with xopen(path, "rb", download_config=download_config) as f:
bytes_ = BytesIO(f.read())
video = VideoReader(bytes_)
else:
video = VideoReader(BytesIO(bytes_))
return video
def flatten(self) -> Union["FeatureType", Dict[str, "FeatureType"]]:
"""If in the decodable state, return the feature itself, otherwise flatten the feature into a dictionary."""
from .features import Value
return (
self
if self.decode
else {
"bytes": Value("binary"),
"path": Value("string"),
}
)
def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArray]) -> pa.StructArray:
"""Cast an Arrow array to the Video arrow storage type.
The Arrow types that can be converted to the Video pyarrow storage type are:
- `pa.string()` - it must contain the "path" data
- `pa.binary()` - it must contain the video bytes
- `pa.struct({"bytes": pa.binary()})`
- `pa.struct({"path": pa.string()})`
- `pa.struct({"bytes": pa.binary(), "path": pa.string()})` - order doesn't matter
- `pa.list(*)` - it must contain the video array data
Args:
storage (`Union[pa.StringArray, pa.StructArray, pa.ListArray]`):
PyArrow array to cast.
Returns:
`pa.StructArray`: Array in the Video arrow storage type, that is
`pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
"""
if pa.types.is_string(storage.type):
bytes_array = pa.array([None] * len(storage), type=pa.binary())
storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"], mask=storage.is_null())
elif pa.types.is_binary(storage.type):
path_array = pa.array([None] * len(storage), type=pa.string())
storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
elif pa.types.is_struct(storage.type):
if storage.type.get_field_index("bytes") >= 0:
bytes_array = storage.field("bytes")
else:
bytes_array = pa.array([None] * len(storage), type=pa.binary())
if storage.type.get_field_index("path") >= 0:
path_array = storage.field("path")
else:
path_array = pa.array([None] * len(storage), type=pa.string())
storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
elif pa.types.is_list(storage.type):
bytes_array = pa.array(
[encode_np_array(np.array(arr))["bytes"] if arr is not None else None for arr in storage.to_pylist()],
type=pa.binary(),
)
path_array = pa.array([None] * len(storage), type=pa.string())
storage = pa.StructArray.from_arrays(
[bytes_array, path_array], ["bytes", "path"], mask=bytes_array.is_null()
)
return array_cast(storage, self.pa_type)
def video_to_bytes(video: "VideoReader") -> bytes:
"""Convert a decord Video object to bytes using native compression if possible"""
raise NotImplementedError()
def encode_decord_video(video: "VideoReader") -> dict:
if hasattr(video, "_hf_encoded"):
return video._hf_encoded
else:
raise NotImplementedError(
"Encoding a decord video is not implemented. "
"Please call `datasets.features.video.patch_decord()` before loading videos to enable this."
)
def encode_np_array(array: np.ndarray) -> dict:
raise NotImplementedError()
# Patching decord a little bit to:
# 1. store the encoded video data {"path": ..., "bytes": ...} in `video._hf_encoded``
# 2. set the decord bridge to numpy/torch/tf/jax using `video._hf_bridge_out` (per video instance) instead of decord.bridge.bridge_out (global)
# This doesn't affect the normal usage of decord.
def _patched_init(self: "VideoReader", uri: Union[str, BytesIO], *args, **kwargs) -> None:
from decord.bridge import bridge_out
if hasattr(uri, "read"):
self._hf_encoded = {"bytes": uri.read(), "path": None}
uri.seek(0)
elif isinstance(uri, str):
self._hf_encoded = {"bytes": None, "path": uri}
self._hf_bridge_out = bridge_out
self._original_init(uri, *args, **kwargs)
def _patched_next(self: "VideoReader", *args, **kwargs):
return self._hf_bridge_out(self._original_next(*args, **kwargs))
def _patched_get_batch(self: "VideoReader", *args, **kwargs):
return self._hf_bridge_out(self._original_get_batch(*args, **kwargs))
def patch_decord():
# We need to import torch first, otherwise later it can cause issues
# e.g. "RuntimeError: random_device could not be read"
# when running `torch.tensor(value).share_memory_()`
# Same for duckdb which crashes on import
if config.TORCH_AVAILABLE:
import torch # noqa
if config.DUCKDB_AVAILABLE:
import duckdb # noqa
import decord.video_reader
from decord import VideoReader
if not hasattr(VideoReader, "_hf_patched"):
decord.video_reader.bridge_out = lambda x: x
VideoReader._original_init = VideoReader.__init__
VideoReader.__init__ = _patched_init
VideoReader._original_next = VideoReader.next
VideoReader.next = _patched_next
VideoReader._original_get_batch = VideoReader.get_batch
VideoReader.get_batch = _patched_get_batch
VideoReader._hf_patched = True
|