Skip to content

Commit f29a8f9

Browse files
tikuma-lsuhsctikuma-lsuhsc
committed
- fixed image output shape
- added squeeze arg to bytes_to_xxx hooks
1 parent 1d4dc16 commit f29a8f9

9 files changed

Lines changed: 41 additions & 20 deletions

File tree

‎src/ffmpegio/audio.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ def _run_read(
6060
)
6161

6262
return rate, plugins.get_hook().bytes_to_audio(
63-
b=out.stdout, dtype=dtype, shape=(ac,)
63+
b=out.stdout, dtype=dtype, shape=(ac,), squeeze=False
6464
)
6565

6666

‎src/ffmpegio/image.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ def _run_read(*args, shape=None, pix_fmt_in=None, s_in=None, show_log=None, **kw
4747
nbytes = utils.get_samplesize(shape, dtype)
4848

4949
return plugins.get_hook().bytes_to_video(
50-
b=out.stdout[-nbytes:], dtype=dtype, shape=shape
50+
b=out.stdout[-nbytes:], dtype=dtype, shape=shape, squeeze=True
5151
)
5252

5353

@@ -192,7 +192,9 @@ def write(url, data, overwrite=None, show_log=None, **options):
192192

193193
ffmpeg_args = configure.empty()
194194
configure.add_url(
195-
ffmpeg_args, "input", *configure.array_to_video_input(1, data=data, **input_options)
195+
ffmpeg_args,
196+
"input",
197+
*configure.array_to_video_input(1, data=data, **input_options),
196198
)
197199
outopts = configure.add_url(ffmpeg_args, "output", url, options)[1][1]
198200
outopts["frames:v"] = 1

‎src/ffmpegio/plugins/hookspecs.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ def audio_bytes(obj: object) -> memoryview:
4949

5050

5151
@hookspec(firstresult=True)
52-
def bytes_to_video(b: bytes, dtype: str, shape: Tuple[int, int, int]) -> object:
52+
def bytes_to_video(b: bytes, dtype: str, shape: Tuple[int, int, int], squeeze: bool) -> object:
5353
"""convert bytes to rawvideo object
5454
5555
:param b: byte data of arbitrary number of video frames
@@ -58,13 +58,15 @@ def bytes_to_video(b: bytes, dtype: str, shape: Tuple[int, int, int]) -> object:
5858
:type dtype: str
5959
:param size: frame dimension in pixels and number of color components (height, width, components)
6060
:type size: Tuple[int, int, int]
61+
:param squeeze: True to remove all the singular dimensions
62+
:type squeeze: bool
6163
:return: python object holding the rawvideo frames
6264
:rtype: object
6365
"""
6466

6567

6668
@hookspec(firstresult=True)
67-
def bytes_to_audio(b: bytes, dtype: str, shape: Tuple[int]) -> object:
69+
def bytes_to_audio(b: bytes, dtype: str, shape: Tuple[int], squeeze: bool) -> object:
6870
"""convert bytes to rawaudio object
6971
7072
:param b: byte data of arbitrary number of video frames
@@ -73,6 +75,8 @@ def bytes_to_audio(b: bytes, dtype: str, shape: Tuple[int]) -> object:
7375
:type dtype: str
7476
:param shape: number of interleaved audio channels (1-element tuple)
7577
:type shape: Tuple[int]
78+
:param squeeze: True to remove all the singular dimensions
79+
:type squeeze: bool
7680
:return: python object to hold the raw audio samples
7781
:rtype: object
7882
"""

‎src/ffmpegio/plugins/rawdata_bytes.py‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,9 @@ def audio_bytes(obj: object) -> memoryview:
5656

5757

5858
@hookimpl
59-
def bytes_to_video(b: bytes, dtype: str, shape: Tuple[int, int, int]) -> object:
59+
def bytes_to_video(
60+
b: bytes, dtype: str, shape: Tuple[int, int, int], squeeze: bool
61+
) -> object:
6062
"""convert bytes to rawvideo object
6163
6264
:param b: byte data of arbitrary number of video frames
@@ -65,19 +67,23 @@ def bytes_to_video(b: bytes, dtype: str, shape: Tuple[int, int, int]) -> object:
6567
:type dtype: str
6668
:param size: frame dimension in pixels and number of color components (height, width, components)
6769
:type size: Tuple[int, int, int]
70+
:param squeeze: True to remove all the singular dimensions
71+
:type squeeze: bool
6872
:return: dict holding the rawvideo frame data
6973
:rtype: dict['buffer':bytes, 'dtype':str, 'shape': Tuple[int,int,int]]
7074
"""
7175

76+
sh = (len(b) // get_samplesize(shape, dtype), *shape)
77+
7278
return {
7379
"buffer": b,
7480
"dtype": dtype,
75-
"shape": (len(b) // get_samplesize(shape, dtype), *shape),
81+
"shape": tuple(((i for i in sh if i != 1))) if squeeze else sh,
7682
}
7783

7884

7985
@hookimpl
80-
def bytes_to_audio(b: bytes, dtype: str, shape: Tuple[int]) -> object:
86+
def bytes_to_audio(b: bytes, dtype: str, shape: Tuple[int], squeeze: bool) -> object:
8187
"""convert bytes to rawaudio object
8288
8389
:param b: byte data of arbitrary number of video frames
@@ -86,12 +92,16 @@ def bytes_to_audio(b: bytes, dtype: str, shape: Tuple[int]) -> object:
8692
:type dtype: str
8793
:param shape: number of interleaved audio channels (1-element tuple)
8894
:type shape: Tuple[int]
95+
:param squeeze: True to remove all the singular dimensions
96+
:type squeeze: bool
8997
:return: dict to hold the raw audio samples
9098
:rtype: dict['buffer':bytes, 'dtype':str, 'shape': Tuple[int]]
9199
"""
92100

101+
sh = (len(b) // get_samplesize(shape, dtype), *shape)
102+
93103
return {
94104
"buffer": b,
95105
"dtype": dtype,
96-
"shape": (len(b) // get_samplesize(shape, dtype), *shape),
106+
"shape": tuple(((i for i in sh if i != 1))) if squeeze else sh,
97107
}

‎src/ffmpegio/streams/SimpleStreams.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ def __init__(
9191

9292
self.samplesize = utils.get_samplesize(self.shape, self.dtype)
9393

94-
self.blocksize = blocksize or max(1024 ** 2 // self.samplesize, 1)
94+
self.blocksize = blocksize or max(1024**2 // self.samplesize, 1)
9595
logging.debug("[reader main] completed init")
9696

9797
def close(self):
@@ -179,7 +179,7 @@ def read(self, n=-1):
179179
logging.debug(f"[reader main] read {len(b)} bytes")
180180
if not len(b):
181181
self._proc.stdout.close()
182-
return self._converter(b=b, shape=self.shape, dtype=self.dtype)
182+
return self._converter(b=b, shape=self.shape, dtype=self.dtype, squeeze=False)
183183

184184
def readinto(self, array):
185185
"""Read bytes into a pre-allocated, writable bytes-like object array and
@@ -428,7 +428,7 @@ def write(self, data):
428428
# the data and start
429429
self._open(data)
430430

431-
logging.debug('[writer main] writing...')
431+
logging.debug("[writer main] writing...")
432432

433433
try:
434434
self._proc.stdin.write(self._viewer(obj=data))
@@ -881,7 +881,7 @@ def filter(self, data, timeout=None):
881881
nread = (int(self.nin * self._out2in) - self.nout) * self._bps_out
882882
y = self._reader.read(-nread, timeout - _time())
883883
self.nout += len(y) // self._bps_out
884-
return self._converter(b=y, dtype=self.dtype, shape=self.shape)
884+
return self._converter(b=y, dtype=self.dtype, shape=self.shape, squeeze=False)
885885

886886
def flush(self, timeout=None):
887887
"""Close the stream input and retrieve the remaining output samples
@@ -900,7 +900,7 @@ def flush(self, timeout=None):
900900
self._proc.wait()
901901
y += self._reader.read_all(None)
902902
self.nout += len(y) // self._bps_out
903-
return self._converter(b=y, dtype=self.dtype, shape=self.shape)
903+
return self._converter(b=y, dtype=self.dtype, shape=self.shape, squeeze=False)
904904

905905

906906
class SimpleVideoFilter(SimpleFilterBase):

‎src/ffmpegio/utils/avi.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -492,7 +492,9 @@ def set_stream_info(hdr):
492492
}
493493

494494
self.streams = {v["index"]: set_stream_info(v) for v in hdr}
495-
self.itemsizes = {v["index"]: get_samplesize(v["shape"], v["dtype"]) for v in hdr}
495+
self.itemsizes = {
496+
v["index"]: get_samplesize(v["shape"], v["dtype"]) for v in hdr
497+
}
496498
self.ready = True
497499

498500
def __next__(self):
@@ -510,7 +512,7 @@ def __iter__(self):
510512
def from_bytes(self, id, b):
511513
info = self.streams[id]
512514
return self.converters[info["type"]](
513-
b=b, dtype=info["dtype"], shape=info["shape"]
515+
b=b, dtype=info["dtype"], shape=info["shape"], squeeze=False
514516
)
515517

516518

‎src/ffmpegio/video.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ def _run_read(
4848
)
4949
if out.returncode:
5050
raise FFmpegError(out.stderr)
51-
return r, plugins.get_hook().bytes_to_video(b=out.stdout, dtype=dtype, shape=shape)
51+
return r, plugins.get_hook().bytes_to_video(b=out.stdout, dtype=dtype, shape=shape, squeeze=False)
5252

5353

5454
def create(

‎tests/test_image.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ def test_create():
1111

1212
# make sure rgb channels are mapped properly
1313
A = image.create("color", c="red", s="100x100", d=1)
14-
assert A["shape"][1:] == (100, 100, 3)
14+
assert A["shape"] == (100, 100, 3)
1515
assert A["dtype"] == ("|u1")
1616
# A = image.create("color", c="green", s="100x100", d=1)
1717
# A = image.create("color", c="blue", s="100x100", d=1)

‎tests/test_plugins.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,17 +8,20 @@ def test_rawdata_bytes():
88
dtype = "|u1"
99
shape = (2, 2, 3)
1010
b = b"\0" * prod(shape)
11-
data = hook.bytes_to_video(b=b, dtype=dtype, shape=shape)
11+
data = hook.bytes_to_video(b=b, dtype=dtype, shape=shape, squeeze=False)
1212
assert data["buffer"] == b
1313
assert data["dtype"] == dtype
1414
assert data["shape"][1:] == shape
1515
assert hook.video_info(obj=data) == (shape, dtype)
1616
assert hook.video_bytes(obj=data) == b
1717

18+
data = hook.bytes_to_video(b=b, dtype=dtype, shape=shape, squeeze=True)
19+
assert data["shape"] == shape
20+
1821
dtype = "<f4"
1922
shape = (2,)
2023
b = b"\0" * (1024 * prod(shape))
21-
data = hook.bytes_to_audio(b=b, dtype=dtype, shape=shape)
24+
data = hook.bytes_to_audio(b=b, dtype=dtype, shape=shape, squeeze=False)
2225
assert data["buffer"] == b
2326
assert data["dtype"] == dtype
2427
assert data["shape"][1:] == shape

0 commit comments

Comments
 (0)