Пример #1
0
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
Пример #2
0
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)
Пример #3
0
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)
Пример #4
0
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)
Пример #5
0
    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
Пример #6
0
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
Пример #7
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()
Пример #8
0
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)
Пример #9
0
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)
Пример #10
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
Пример #11
0
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
Пример #12
0
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
Пример #13
0
def test_duration(video_filename):
    allocator = TorchImageAllocator('cpu', torch.uint8)
    fr = pyfffr.TvFFFrameReader(allocator, video_filename, -1)
    assert fr.get_duration() == 2.0
Пример #14
0
def test_get_height(video_filename):
    allocator = TorchImageAllocator('cpu', torch.uint8)
    fr = pyfffr.TvFFFrameReader(allocator, video_filename, -1)
    assert fr.get_height() == 720
Пример #15
0
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