File size: 1,205 Bytes
aa7758f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | from __future__ import annotations
import unittest
import numpy as np
import torch
from qprefer_reward.processing import frame_to_pil, uniform_sample_frames
class ProcessingTest(unittest.TestCase):
def test_minus_one_one_tensor_to_rgb(self) -> None:
frame = torch.tensor(
[
[[-1.0, 1.0], [-1.0, 1.0]],
[[-1.0, 1.0], [-1.0, 1.0]],
[[-1.0, 1.0], [-1.0, 1.0]],
],
dtype=torch.float32,
)
image = frame_to_pil(frame, value_range="minus_one_one")
pixels = np.asarray(image)
self.assertEqual(image.mode, "RGB")
self.assertEqual(pixels.shape, (2, 2, 3))
self.assertEqual(int(pixels.min()), 0)
self.assertEqual(int(pixels.max()), 255)
def test_uniform_sampler_repeats_short_video(self) -> None:
video = np.zeros((2, 4, 4, 3), dtype=np.uint8)
video[1] = 255
frames = uniform_sample_frames(video, num_frames=8)
self.assertEqual(len(frames), 8)
self.assertEqual(int(np.asarray(frames[0]).mean()), 0)
self.assertEqual(int(np.asarray(frames[-1]).mean()), 255)
if __name__ == "__main__":
unittest.main()
|