class TestGetBatchPytorch(unittest.TestCase):
def setUp(self):
self.video = Video('./tests/test_video.mp4', data_container='pytorch')
def test_pytorch_shape(self):
batch = self.video.get_batch([0, 1])
self.assertTrue(torch.is_tensor(batch))
self.assertEqual(batch.size(), (2, 456, 256, 3))
class TestGetBatch(unittest.TestCase):
def setUp(self):
self.video = Video('./tests/test_video.mp4')
def test_numpy_shape(self):
batch = self.video.get_batch([0, 1])
self.assertIsInstance(batch, np.ndarray)
self.assertEqual(batch.shape, (2, 456, 256, 3))
def test_pytorch_comp_with_ffmpeg(self):
video = Video('./tests/test_video.mp4', data_container='pytorch')
frames = list(range(2))
batch = video.get_batch(frames)
expected = self.get_truth(batch.shape, frames)
expected = torch.from_numpy(expected.copy())
torch.testing.assert_allclose(batch, expected, 0, 0)
class TestGetBatchBasic(unittest.TestCase):
def setUp(self):
self.video = Video('./tests/test_video.mp4', data_container=None)
self.batch = self.video.get_batch([0, 1])
def test_type(self):
self.assertRegex(repr(self.batch),
'<capsule object "dltensor" at 0x[0-9a-f]+>')
def test_len(self):
self.assertEqual(self.video.num_frames(), 300)
self.assertEqual(len(self.video), 300)
def test_out_of_range(self):
with self.assertRaises(IndexError):
self.video.get_batch([9999])
def test_fps(self):
fps = self.video.average_frame_rate()
self.assertIsInstance(fps, fractions.Fraction)
self.assertEqual(fps, fractions.Fraction(30000, 1001))
def test_comp_with_ffmpeg(self):
video = Video('./tests/test_video.mp4')
frameIndices = [
range(1),
range(2),
range(3),
range(8),
range(40),
range(3, 18),
list(range(3, 18)) + list(range(66, 88)),
[1, 1, 1],
[69, 55, 22, 80, 32],
]
for frames in frameIndices:
frames = list(frames)
with self.subTest(frame_indices=frames):
batch = video.get_batch(frames)
expected = self.get_truth(batch.shape, frames)
for i, f in enumerate(frames):
numpy.testing.assert_array_equal(
batch[i], expected[i],
f'Frame {f} mismatch ({i}th frame in batch)')
def test_keep_awake(self):
video = Video('./tests/test_video.mp4')
with video.keep_awake():
video.get_batch([0, 1])
self.assertFalse(video.is_sleeping())
self.assertTrue(video.is_sleeping())
def test_sleeping_after_read(self):
video = Video('./tests/test_video.mp4')
video.get_batch([0, 1])
self.assertTrue(video.is_sleeping())