def test_to_uint8(self):
self.assert_all_equal(
media.to_uint8(np.array([0, 1, 128, 254, 255], dtype=np.uint8)),
[0, 1, 128, 254, 255])
self.assert_all_close(
media.to_uint8([-0.2, 0.0, 0.1, 0.5, 0.9, 1.0, 1.1]), [
0, 0,
int(0.1 * 255 + 0.5),
int(0.5 * 255 + 0.5),
int(0.9 * 255 + 0.5), 255, 255
])
def test_video_streaming_write_read_roundtrip(self):
shape = (62, 744)
num_images = 20
fps = 120
bps = 400_000
with tempfile.TemporaryDirectory() as directory_name:
filename = os.path.join(directory_name, 'test.mp4')
images = []
with media.VideoWriter(filename, shape, fps=fps,
bps=bps) as writer:
for image in media.moving_circle(shape, num_images):
image_uint8 = media.to_uint8(image)
writer.add_image(image_uint8)
images.append(image_uint8)
with media.VideoReader(filename) as reader:
self.assertEqual(reader.num_images, num_images)
self.assert_all_equal(reader.shape, shape)
self.assertEqual(reader.fps, fps)
self.assertBetween(reader.bps, 100_000, 500_000)
self.assertEqual(reader.metadata.num_images, reader.num_images)
self.assertEqual(reader.metadata.shape, reader.shape)
self.assertEqual(reader.metadata.fps, reader.fps)
self.assertEqual(reader.metadata.bps, reader.bps)
for index, new_image in enumerate(reader):
self._check_similar(images[index], new_image, 7.0,
f'index={index}')
def test_rgb_ycbcr_roundtrip(self):
image = np.array(
[[0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255], [255, 255, 0],
[0, 255, 255], [255, 0, 255], [255, 255, 255], [128, 128, 128]],
dtype=np.uint8)
new = media.to_uint8(media.rgb_from_ycbcr(media.ycbcr_from_rgb(image)))
self.assert_all_close(image, new, atol=1)
def test_video_non_streaming_write_read_roundtrip(self, use_generator):
shape = (240, 320)
num_images = 10
fps = 40
qp = 20
original_video = media.to_uint8(media.moving_circle(shape, num_images))
video = (
image
for image in original_video) if use_generator else original_video
with tempfile.TemporaryDirectory() as directory_name:
filename = os.path.join(directory_name, 'test.mp4')
media.write_video(filename, video, fps=fps, qp=qp)
new_video = media.read_video(filename)
self._check_similar(original_video, new_video, 3.0)
def test_video_streaming_read_write(self):
shape = (400, 400)
num_images = 4
fps = 60
bps = 40_000_000
video = media.to_uint8(media.moving_circle(shape, num_images))
with tempfile.TemporaryDirectory() as directory_name:
filename1 = os.path.join(directory_name, 'test1.mp4')
filename2 = os.path.join(directory_name, 'test2.mp4')
media.write_video(filename1, video, fps=fps, bps=bps)
with media.VideoReader(filename1) as reader:
with media.VideoWriter(filename2,
reader.shape,
fps=reader.fps,
bps=reader.bps,
encoded_format='yuv420p') as writer:
for image in reader:
writer.add_image(image)
new_video = media.read_video(filename2)
self._check_similar(video, new_video, 3.0)