def test_allocate_frame(device, dtype):
dtype = getattr(torch, dtype)
mm = TorchImageAllocator(device, dtype)
address = mm.allocate_frame(16, 9, 16, 32)
assert address
tensor = mm.get_frame_tensor(address)
assert tensor.shape == (3, 9, 16)
assert tensor.device == torch.device(device)
assert tensor.dtype == dtype
assert tensor.data_ptr() == address
def test_read_frame_uint8(device, video_filename, first_frame_image):
allocator = TorchImageAllocator(device, torch.uint8)
gpu_index = allocator.device.index if allocator.device.type == 'cuda' else -1
fr = pyfffr.TvFFFrameReader(allocator, video_filename, gpu_index)
ptr = fr.read_frame()
assert ptr is not None
rgb_frame = allocator.get_frame_tensor(int(ptr))
assert rgb_frame.shape == (3, 720, 1280)
actual = PIL.Image.fromarray(rgb_frame.cpu().permute(1, 2, 0).numpy(),
'RGB')
assert_allclose(actual, first_frame_image, atol=50)
def test_allocation_alignment(device):
dtype = torch.float32
allocator = TorchImageAllocator(device, dtype)
addr = allocator.allocate_frame(1, 1, 4, 32)
tensor = allocator.get_frame_tensor(addr)
storage = tensor.storage()
# Not using Storage#fill_ because of https://github.com/pytorch/pytorch/issues/22383
for i in range(len(storage)):
storage[i] = 0
tensor.fill_(1)
for i in range(3):
index = tensor.storage_offset() + i * 8
assert storage[index] == 1
assert (len(storage) == 32)
def test_read_frame_float32(device, video_filename, first_frame_image):
if device == 'cpu':
pytest.skip(
'Software decoding does not support float32 output directly')
allocator = TorchImageAllocator(device, torch.float32)
gpu_index = allocator.device.index if allocator.device.type == 'cuda' else -1
fr = pyfffr.TvFFFrameReader(allocator, video_filename, gpu_index)
ptr = fr.read_frame()
assert ptr is not None
rgb_frame = allocator.get_frame_tensor(int(ptr))
assert rgb_frame.shape == (3, 720, 1280)
actual = PIL.Image.fromarray(
(rgb_frame.cpu() * 255).byte().permute(1, 2, 0).numpy(), 'RGB')
assert_allclose(actual, first_frame_image, atol=50)
def __init__(self,
filename,
device,
dtype,
*,
seek_threshold=0,
out_width=0,
out_height=0,
buffer_length=8):
super().__init__(filename, device, dtype, seek_threshold, out_width,
out_height)
self.lock = RLock()
allocator_dtype = self.dtype
# The FFFR backend does not currently support direct conversion to float32 for software
# decoding, so we will read as uint8 and do the data type conversion afterwards.
if self.device.type == 'cpu' and self.dtype != torch.uint8:
allocator_dtype = torch.uint8
image_allocator = TorchImageAllocator(self.device, allocator_dtype)
device_index = self.device.index if self.device.type == 'cuda' else -1
frame_reader = pyfffr.TvFFFrameReader(image_allocator, self.filename,
device_index, out_width,
out_height, self.seek_threshold,
buffer_length)
# We need to hold a reference to image_allocator for at least as long as the
# TvFFFrameReader that uses it is around, since we retain ownership of image_allocator.
setattr(frame_reader, '__image_allocator_ref', image_allocator)
self.image_allocator = image_allocator
self.frame_reader = frame_reader
self._at_eof = False
def test_device_works_after_reading_frame(device, video_filename):
allocator = TorchImageAllocator(device, torch.uint8)
cuda_tensor = torch.zeros(1, device=allocator.device)
assert cuda_tensor.item() == 0
gpu_index = allocator.device.index if allocator.device.type == 'cuda' else -1
fr = pyfffr.TvFFFrameReader(allocator, video_filename, gpu_index)
fr.read_frame()
assert cuda_tensor.item() == 0
def test_two_decoders(video_filename):
allocator = TorchImageAllocator('cuda:0', torch.uint8)
fr1 = pyfffr.TvFFFrameReader(allocator, video_filename,
allocator.device.index)
fr1.read_frame()
fr2 = pyfffr.TvFFFrameReader(allocator, video_filename,
allocator.device.index)
fr2.read_frame()
def test_read_frames_by_index(device, video_filename, mid_frame_image):
allocator = TorchImageAllocator(device, torch.uint8)
gpu_index = allocator.device.index if allocator.device.type == 'cuda' else -1
fr = pyfffr.TvFFFrameReader(allocator, video_filename, gpu_index)
fr.seek(0.4)
frame_offsets = torch.tensor([0, 10, 25, 30],
device='cpu',
dtype=torch.int64)
ptrs = torch.zeros(frame_offsets.shape, device='cpu', dtype=torch.int64)
res = fr.read_frames_by_index(frame_offsets.data_ptr(), len(frame_offsets),
ptrs.data_ptr())
assert res == len(frame_offsets)
rgb_frame = allocator.get_frame_tensor(int(ptrs[2]))
assert rgb_frame.shape == (3, 720, 1280)
actual = PIL.Image.fromarray(rgb_frame.cpu().permute(1, 2, 0).numpy(),
'RGB')
assert_allclose(actual, mid_frame_image, atol=50)
def test_seek_out_of_bounds(video_filename):
allocator = TorchImageAllocator('cpu', torch.uint8)
fr = pyfffr.TvFFFrameReader(allocator, video_filename, -1)
# Try seeking to a time after the end of the video.
with pytest.raises(RuntimeError):
fr.seek(2.0)
# Try seeking to a time before the start of the video.
with pytest.raises(RuntimeError):
fr.seek(-1.0)
# Check that seeking within bounds is still OK.
fr.seek(1.0)
def test_free_frame(device):
mm = TorchImageAllocator(device, torch.uint8)
addr1 = mm.allocate_frame(16, 9, 16, 32)
addr2 = mm.allocate_frame(16, 9, 16, 32)
assert len(mm.tensors) == 2
mm.free_frame(addr1)
assert len(mm.tensors) == 1
assert mm.get_frame_tensor(addr2) is not None
def test_get_device_index():
allocator = TorchImageAllocator('cuda', torch.float32)
assert allocator.get_device_index() >= 0
allocator = TorchImageAllocator('cpu', torch.float32)
assert allocator.get_device_index() == -1
def test_get_data_type(device, dtype, data_type_enum_value):
dtype = getattr(torch, dtype)
allocator = TorchImageAllocator(device, dtype)
assert allocator.get_data_type() == data_type_enum_value
def test_duration(video_filename):
allocator = TorchImageAllocator('cpu', torch.uint8)
fr = pyfffr.TvFFFrameReader(allocator, video_filename, -1)
assert fr.get_duration() == 2.0
def test_get_height(video_filename):
allocator = TorchImageAllocator('cpu', torch.uint8)
fr = pyfffr.TvFFFrameReader(allocator, video_filename, -1)
assert fr.get_height() == 720
def test_n_frames(video_filename):
allocator = TorchImageAllocator('cpu', torch.uint8)
fr = pyfffr.TvFFFrameReader(allocator, video_filename, -1)
assert fr.get_number_of_frames() == 50
