| 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() |
|
|
