-
Notifications
You must be signed in to change notification settings - Fork 9
Expand file tree
/
Copy pathrawdata_numpy.py
More file actions
128 lines (102 loc) · 3.32 KB
/
rawdata_numpy.py
File metadata and controls
128 lines (102 loc) · 3.32 KB
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
"""ffmpegio plugin to use `numpy.ndarray` objects for media data I/O"""
import numpy as np
from pluggy import HookimplMarker
from typing import Tuple
from numpy.typing import ArrayLike
hookimpl = HookimplMarker("ffmpegio")
__all__ = [
"video_info",
"audio_info",
"video_bytes",
"audio_bytes",
"bytes_to_video",
"bytes_to_audio",
]
@hookimpl
def video_info(obj: ArrayLike) -> Tuple[Tuple[int, int, int], str]:
"""get video frame info
:param obj: video frame data with arbitrary number of frames
:type obj: ArrayLike
:return: shape (height,width,components) and data type str
:rtype: Tuple[Tuple[int, int, int], str]
"""
try:
return obj.shape[-3:] if obj.ndim != 2 else [*obj.shape, 1], obj.dtype.str
except:
return None
@hookimpl
def audio_info(obj: ArrayLike) -> Tuple[int, str]:
"""get audio sample info
:param obj: column-wise audio data with arbitrary number of samples
:type obj: ArrayLike
:return: number of channels and sample data type in data type str
:rtype: Tuple[Tuple[int], str]
"""
try:
return obj.shape[-1:] if obj.ndim > 1 else [1], obj.dtype.str
except:
return None
@hookimpl
def video_bytes(obj: ArrayLike) -> memoryview:
"""return bytes-like object of rawvideo NumPy array
:param obj: video frame data with arbitrary number of frames
:type obj: ArrayLike
:return: memoryview of video frames
:rtype: memoryview
"""
try:
return memoryview(np.ascontiguousarray(obj, obj.dtype))
except:
return None
@hookimpl
def audio_bytes(obj: ArrayLike) -> memoryview:
"""return bytes-like object of rawaudio NumPy array
:param obj: column-wise audio data with arbitrary number of samples
:type obj: ArrayLike
:return: memoryview of audio samples
:rtype: memoryview
"""
try:
return memoryview(np.ascontiguousarray(obj, obj.dtype))
except:
return None
@hookimpl
def bytes_to_video(
b: bytes, dtype: str, shape: Tuple[int, int, int], squeeze: bool
) -> ArrayLike:
"""convert bytes to rawvideo NumPy array
:param b: byte data of arbitrary number of video frames
:type b: bytes
:param dtype: data type string (e.g., '|u1', '<f4')
:type dtype: str
:param size: frame dimension in pixels and number of color components (height, width, components)
:type size: Tuple[int, int, int]
:param squeeze: True to remove all the singular dimensions
:type squeeze: bool
:return: rawvideo frames
:rtype: ArrayLike
"""
try:
x = np.frombuffer(b, dtype).reshape(-1, *shape)
return x.squeeze() if squeeze else x
except:
return None
@hookimpl
def bytes_to_audio(b: bytes, dtype: str, shape: Tuple[int], squeeze: bool) -> ArrayLike:
"""convert bytes to rawaudio NumPy array
:param b: byte data of arbitrary number of video frames
:type b: bytes
:param dtype: data type string (e.g., '<s2', '<f4')
:type dtype: str
:param shape: number of audio channels
:type shape: Tuple[int]
:param squeeze: True to remove all the singular dimensions
:type squeeze: bool
:return: raw audio samples
:rtype: ArrayLike
"""
try:
x = np.frombuffer(b, dtype).reshape(-1, *shape)
return x.squeeze() if squeeze else x
except:
return None