Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 34 additions & 15 deletions av/video/frame.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ cdef object _cinit_bypass_sentinel
supported_np_pix_fmts = {
"abgr", "argb", "bayer_bggr16be", "bayer_bggr16le", "bayer_bggr8", "bayer_gbrg16be",
"bayer_gbrg16le", "bayer_gbrg8", "bayer_grbg16be", "bayer_grbg16le", "bayer_grbg8",
"bayer_rggb16be", "bayer_rggb16le", "bayer_rggb8", "bgr24", "bgr8", "bgra",
"bayer_rggb16be", "bayer_rggb16le", "bayer_rggb8", "bgr24", "bgr48be", "bgr48le", "bgr8", "bgra", "bgra64be", "bgra64le",
"gbrapf32be", "gbrapf32le", "gbrp", "gbrp10be", "gbrp10le", "gbrp12be", "gbrp12le",
"gbrp14be", "gbrp14le", "gbrp16be", "gbrp16le", "gbrpf32be", "gbrpf32le", "gray",
"gray16be", "gray16le", "gray8", "grayf32be", "grayf32le", "nv12", "pal8", "rgb24",
Expand Down Expand Up @@ -342,6 +342,8 @@ cdef class VideoFrame(Frame):
"bayer_rggb16le": (2, "uint16"),
"bayer_rggb16be": (2, "uint16"),
"bgr24": (3, "uint8"),
"bgr48be": (6, "uint16"),
"bgr48le": (6, "uint16"),
"bgr8": (1, "uint8"),
"bgra": (4, "uint8"),
"gbrapf32be": (4, "float32"),
Expand Down Expand Up @@ -370,6 +372,8 @@ cdef class VideoFrame(Frame):
"rgba": (4, "uint8"),
"rgba64be": (8, "uint16"),
"rgba64le": (8, "uint16"),
"bgra64be": (8, "uint16"),
"bgra64le": (8, "uint16"),
"yuv444p": (1, "uint8"),
"yuv444p16be": (2, "uint16"),
"yuv444p16le": (2, "uint16"),
Expand Down Expand Up @@ -467,43 +471,58 @@ cdef class VideoFrame(Frame):
if not width:
width = array.shape[1]

if format in ("rgb24", "bgr24"):
if format in {"rgb24", "bgr24"}:
check_ndarray(array, "uint8", 3)
check_ndarray_shape(array, array.shape[2] == 3)
if array.strides[1:] != (3, 1):
raise ValueError("provided array does not have C_CONTIGUOUS rows")
linesizes = (array.strides[0], )
elif format in {"rgb48le", "rgb48be", "bgr48le", "bgr48be"}:
check_ndarray(array, "uint16", 3)
check_ndarray_shape(array, array.shape[2] == 3)
if array.strides[1:] != (6, 2):
raise ValueError("provided array does not have C_CONTIGUOUS rows")
linesizes = (array.strides[0], )
elif format in ("rgba", "bgra"):
elif format in {"rgba", "bgra", "argb", "abgr"}:
check_ndarray(array, "uint8", 3)
check_ndarray_shape(array, array.shape[2] == 4)
if array.strides[1:] != (4, 1):
raise ValueError("provided array does not have C_CONTIGUOUS rows")
linesizes = (array.strides[0], )
elif format in ("gray", "gray8", "rgb8", "bgr8"):
elif format in {"rgba64le", "rgba64be", "bgra64le", "bgra64be"}:
check_ndarray(array, "uint16", 3)
check_ndarray_shape(array, array.shape[2] == 4)
if array.strides[1:] != (8, 2):
raise ValueError("provided array does not have C_CONTIGUOUS rows")
linesizes = (array.strides[0], )
elif format in {"gray", "gray8", "rgb8", "bgr8","bayer_bggr8", "bayer_rggb8", "bayer_gbrg8", "bayer_grbg8"}:
check_ndarray(array, "uint8", 2)
if array.strides[1] != 1:
raise ValueError("provided array does not have C_CONTIGUOUS rows")
linesizes = (array.strides[0], )
elif format in ("yuv420p", "yuvj420p", "nv12"):
elif format in {"gray16le", "gray16be", "bayer_rggb16le", "bayer_gbrg16le", "bayer_grbg16le","bayer_bggr16be", "bayer_rggb16be", "bayer_gbrg16be", "bayer_grbg16be"}:
check_ndarray(array, "uint16", 2)
if array.strides[1] != 2:
raise ValueError("provided array does not have C_CONTIGUOUS rows")
linesizes = (array.strides[0], )
elif format in {"grayf32le", "grayf32be"}:
check_ndarray(array, "float32", 2)
if array.strides[1] != 4:
raise ValueError("provided array does not have C_CONTIGUOUS rows")
linesizes = (array.strides[0], )
elif format in {"yuv420p", "yuvj420p", "nv12"}:
check_ndarray(array, "uint8", 2)
check_ndarray_shape(array, array.shape[0] % 3 == 0)
check_ndarray_shape(array, array.shape[1] % 2 == 0)
height = height // 6 * 4
if array.strides[1] != 1:
raise ValueError("provided array does not have C_CONTIGUOUS rows")
if format in ("yuv420p", "yuvj420p"):
if format in {"yuv420p", "yuvj420p"}:
# For YUV420 planar formats, the UV plane stride is always half the Y stride.
linesizes = (array.strides[0], array.strides[0] // 2, array.strides[0] // 2)
else:
# Planes where U and V are interleaved have the same stride as Y.
linesizes = (array.strides[0], array.strides[0])
elif format in {"bayer_bggr8", "bayer_rggb8", "bayer_gbrg8", "bayer_grbg8","bayer_bggr16le", "bayer_rggb16le", "bayer_gbrg16le", "bayer_grbg16le","bayer_bggr16be", "bayer_rggb16be", "bayer_gbrg16be", "bayer_grbg16be"}:
check_ndarray(array, "uint8" if format.endswith("8") else "uint16", 2)

if array.strides[1] != (1 if format.endswith("8") else 2):
raise ValueError("provided array does not have C_CONTIGUOUS rows")

linesizes = (array.strides[0],)
else:
raise ValueError(f"Conversion from numpy array with format `{format}` is not yet supported")

Expand Down Expand Up @@ -717,13 +736,13 @@ cdef class VideoFrame(Frame):
elif format in {"argb", "rgba", "abgr", "bgra"}:
check_ndarray(array, "uint8", 3)
check_ndarray_shape(array, array.shape[2] == 4)
elif format in {"rgb48be", "rgb48le"}:
elif format in {"rgb48be", "rgb48le","bgr48be", "bgr48le"}:
check_ndarray(array, "uint16", 3)
check_ndarray_shape(array, array.shape[2] == 3)
frame = VideoFrame(array.shape[1], array.shape[0], format)
copy_array_to_plane(byteswap_array(array, format.endswith("be")), frame.planes[0], 6)
return frame
elif format in {"rgba64be", "rgba64le"}:
elif format in {"rgba64be", "rgba64le", "bgra64be", "bgra64le"}:
check_ndarray(array, "uint16", 3)
check_ndarray_shape(array, array.shape[2] == 4)
frame = VideoFrame(array.shape[1], array.shape[0], format)
Expand Down
103 changes: 103 additions & 0 deletions tests/test_videoframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -285,6 +285,35 @@ def test_ndarray_rgba_align() -> None:
assertNdarraysEqual(frame.to_ndarray(), array)


def test_ndarray_bayer8() -> None:
array = numpy.random.randint(0, 256, size=(480, 640), dtype=numpy.uint8)
for format in ("bayer_bggr8", "bayer_gbrg8", "bayer_grbg8", "bayer_rggb8"):
frame = VideoFrame.from_ndarray(array, format=format)
assert format in av.video.frame.supported_np_pix_fmts
assert frame.width == 640 and frame.height == 480
assert frame.format.name == format
assertNdarraysEqual(frame.to_ndarray(), array)


def test_ndarray_bayer16() -> None:
array = numpy.random.randint(0, 65536, size=(480, 640), dtype=numpy.uint16)
for format in (
"bayer_bggr16be",
"bayer_bggr16le",
"bayer_gbrg16be",
"bayer_gbrg16le",
"bayer_grbg16be",
"bayer_grbg16le",
"bayer_rggb16be",
"bayer_rggb16le",
):
frame = VideoFrame.from_ndarray(array, format=format)
assert format in av.video.frame.supported_np_pix_fmts
assert frame.width == 640 and frame.height == 480
assert frame.format.name == format
assertNdarraysEqual(frame.to_ndarray(), array)


def test_ndarray_gbrp() -> None:
array = numpy.random.randint(0, 256, size=(480, 640, 3), dtype=numpy.uint8)
frame = VideoFrame.from_ndarray(array, format="gbrp")
Expand Down Expand Up @@ -571,6 +600,17 @@ def test_ndarray_rgb48be() -> None:
assertPixelValue16(frame.planes[0], array[0][0][0], "big")


def test_ndarray_bgr48be() -> None:
array = numpy.random.randint(0, 65536, size=(480, 640, 3), dtype=numpy.uint16)
frame = VideoFrame.from_ndarray(array, format="bgr48be")
assert frame.width == 640 and frame.height == 480
assert frame.format.name == "bgr48be"
assertNdarraysEqual(frame.to_ndarray(), array)

# check endianness by examining blue value of first pixel
assertPixelValue16(frame.planes[0], array[0][0][0], "big")


def test_ndarray_rgb48le() -> None:
array = numpy.random.randint(0, 65536, size=(480, 640, 3), dtype=numpy.uint16)
frame = VideoFrame.from_ndarray(array, format="rgb48le")
Expand All @@ -582,6 +622,17 @@ def test_ndarray_rgb48le() -> None:
assertPixelValue16(frame.planes[0], array[0][0][0], "little")


def test_ndarray_bgr48le() -> None:
array = numpy.random.randint(0, 65536, size=(480, 640, 3), dtype=numpy.uint16)
frame = VideoFrame.from_ndarray(array, format="bgr48le")
assert frame.width == 640 and frame.height == 480
assert frame.format.name == "bgr48le"
assertNdarraysEqual(frame.to_ndarray(), array)

# check endianness by examining blue value of first pixel
assertPixelValue16(frame.planes[0], array[0][0][0], "little")


def test_ndarray_rgb48le_align() -> None:
array = numpy.random.randint(0, 65536, size=(238, 318, 3), dtype=numpy.uint16)
frame = VideoFrame.from_ndarray(array, format="rgb48le")
Expand All @@ -593,6 +644,17 @@ def test_ndarray_rgb48le_align() -> None:
assertPixelValue16(frame.planes[0], array[0][0][0], "little")


def test_ndarray_bgr48le_align() -> None:
array = numpy.random.randint(0, 65536, size=(238, 318, 3), dtype=numpy.uint16)
frame = VideoFrame.from_ndarray(array, format="bgr48le")
assert frame.width == 318 and frame.height == 238
assert frame.format.name == "bgr48le"
assertNdarraysEqual(frame.to_ndarray(), array)

# check endianness by examining blue value of first pixel
assertPixelValue16(frame.planes[0], array[0][0][0], "little")


def test_ndarray_rgba64be() -> None:
array = numpy.random.randint(0, 65536, size=(480, 640, 4), dtype=numpy.uint16)
frame = VideoFrame.from_ndarray(array, format="rgba64be")
Expand All @@ -604,6 +666,17 @@ def test_ndarray_rgba64be() -> None:
assertPixelValue16(frame.planes[0], array[0][0][0], "big")


def test_ndarray_bgra64be() -> None:
array = numpy.random.randint(0, 65536, size=(480, 640, 4), dtype=numpy.uint16)
frame = VideoFrame.from_ndarray(array, format="bgra64be")
assert frame.width == 640 and frame.height == 480
assert frame.format.name == "bgra64be"
assertNdarraysEqual(frame.to_ndarray(), array)

# check endianness by examining blue value of first pixel
assertPixelValue16(frame.planes[0], array[0][0][0], "big")


def test_ndarray_rgba64le() -> None:
array = numpy.random.randint(0, 65536, size=(480, 640, 4), dtype=numpy.uint16)
frame = VideoFrame.from_ndarray(array, format="rgba64le")
Expand All @@ -615,6 +688,17 @@ def test_ndarray_rgba64le() -> None:
assertPixelValue16(frame.planes[0], array[0][0][0], "little")


def test_ndarray_bgra64le() -> None:
array = numpy.random.randint(0, 65536, size=(480, 640, 4), dtype=numpy.uint16)
frame = VideoFrame.from_ndarray(array, format="bgra64le")
assert frame.width == 640 and frame.height == 480
assert frame.format.name == "bgra64le"
assertNdarraysEqual(frame.to_ndarray(), array)

# check endianness by examining blue value of first pixel
assertPixelValue16(frame.planes[0], array[0][0][0], "little")


def test_ndarray_rgb8() -> None:
array = numpy.random.randint(0, 256, size=(480, 640), dtype=numpy.uint8)
frame = VideoFrame.from_ndarray(array, format="rgb8")
Expand Down Expand Up @@ -805,6 +889,25 @@ def test_shares_memory_rgba() -> None:
assertNdarraysEqual(frame.to_ndarray(), array)


def test_shares_memory_bayer8() -> None:
for format in ("bayer_rggb8", "bayer_bggr8", "bayer_grbg8", "bayer_gbrg8"):
array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8)
frame = VideoFrame.from_numpy_buffer(array, format)
assertNdarraysEqual(frame.to_ndarray(), array)

array[...] = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8)
assertNdarraysEqual(frame.to_ndarray(), array)

array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8)
array = array[:, :300]
assert not array.data.c_contiguous
frame = VideoFrame.from_numpy_buffer(array, format)
assertNdarraysEqual(frame.to_ndarray(), array)

array[...] = numpy.random.randint(0, 256, size=array.shape, dtype=numpy.uint8)
assertNdarraysEqual(frame.to_ndarray(), array)


def test_shares_memory_yuv420p() -> None:
array = numpy.random.randint(0, 256, size=(512 * 6 // 4, 256), dtype=numpy.uint8)
frame = VideoFrame.from_numpy_buffer(array, "yuv420p")
Expand Down
Loading