|
2 | 2 |
|
3 | 3 | """Tests for io_utils extensionless video file handling (D99228861).""" |
4 | 4 |
|
| 5 | +import os |
5 | 6 | import tempfile |
6 | 7 | import unittest |
7 | 8 | from unittest.mock import MagicMock, patch |
8 | 9 |
|
| 10 | +import numpy as np |
| 11 | +import torch |
| 12 | + |
9 | 13 | from sam3.model.io_utils import load_video_frames |
10 | 14 |
|
11 | 15 |
|
@@ -107,6 +111,89 @@ def test_dummy_video_pattern(self) -> None: |
107 | 111 | self.assertEqual(h, 480) |
108 | 112 | self.assertEqual(w, 640) |
109 | 113 |
|
| 114 | + def test_cv2_video_file_loader_scales_before_normalization(self) -> None: |
| 115 | + """OpenCV video loading should match normalized decoded uint8 frames.""" |
| 116 | + try: |
| 117 | + import cv2 |
| 118 | + except ImportError as exc: |
| 119 | + self.skipTest(f"OpenCV is required for this test: {exc}") |
| 120 | + |
| 121 | + image_size = 8 |
| 122 | + source_height = 6 |
| 123 | + source_width = 10 |
| 124 | + img_mean = (0.5, 0.25, 0.75) |
| 125 | + img_std = (0.5, 0.25, 0.25) |
| 126 | + yy, xx = np.indices((source_height, source_width), dtype=np.uint16) |
| 127 | + frames_rgb = [ |
| 128 | + np.stack( |
| 129 | + ( |
| 130 | + (xx * 23 + yy * 7) % 256, |
| 131 | + (xx * 11 + yy * 17 + 3) % 256, |
| 132 | + (xx * 5 + yy * 29 + 9) % 256, |
| 133 | + ), |
| 134 | + axis=-1, |
| 135 | + ).astype(np.uint8), |
| 136 | + np.stack( |
| 137 | + ( |
| 138 | + (xx * 13 + yy * 19 + 31) % 256, |
| 139 | + (xx * 3 + yy * 41 + 47) % 256, |
| 140 | + (xx * 37 + yy * 2 + 61) % 256, |
| 141 | + ), |
| 142 | + axis=-1, |
| 143 | + ).astype(np.uint8), |
| 144 | + ] |
| 145 | + |
| 146 | + with tempfile.TemporaryDirectory() as tmpdir: |
| 147 | + video_path = os.path.join(tmpdir, "tiny.avi") |
| 148 | + writer = cv2.VideoWriter( |
| 149 | + video_path, |
| 150 | + cv2.VideoWriter_fourcc(*"MJPG"), |
| 151 | + 2.0, |
| 152 | + (source_width, source_height), |
| 153 | + ) |
| 154 | + if not writer.isOpened(): |
| 155 | + self.skipTest("OpenCV could not create a temporary MJPG video") |
| 156 | + |
| 157 | + for frame_rgb in frames_rgb: |
| 158 | + writer.write(cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)) |
| 159 | + writer.release() |
| 160 | + |
| 161 | + decoded_frames = [] |
| 162 | + cap = cv2.VideoCapture(video_path) |
| 163 | + while True: |
| 164 | + ret, frame_bgr = cap.read() |
| 165 | + if not ret: |
| 166 | + break |
| 167 | + frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) |
| 168 | + decoded_frames.append( |
| 169 | + cv2.resize( |
| 170 | + frame_rgb, |
| 171 | + (image_size, image_size), |
| 172 | + interpolation=cv2.INTER_CUBIC, |
| 173 | + ) |
| 174 | + ) |
| 175 | + cap.release() |
| 176 | + self.assertEqual(len(decoded_frames), len(frames_rgb)) |
| 177 | + |
| 178 | + expected = torch.from_numpy(np.stack(decoded_frames, axis=0)) |
| 179 | + expected = expected.permute(0, 3, 1, 2).to(dtype=torch.float16) |
| 180 | + expected /= 255.0 |
| 181 | + expected -= torch.tensor(img_mean, dtype=torch.float16).view(1, 3, 1, 1) |
| 182 | + expected /= torch.tensor(img_std, dtype=torch.float16).view(1, 3, 1, 1) |
| 183 | + |
| 184 | + frames, height, width = load_video_frames( |
| 185 | + video_path=video_path, |
| 186 | + image_size=image_size, |
| 187 | + offload_video_to_cpu=True, |
| 188 | + img_mean=img_mean, |
| 189 | + img_std=img_std, |
| 190 | + video_loader_type="cv2", |
| 191 | + ) |
| 192 | + |
| 193 | + self.assertEqual((height, width), (source_height, source_width)) |
| 194 | + self.assertEqual(frames.dtype, torch.float16) |
| 195 | + torch.testing.assert_close(frames, expected, rtol=0, atol=1e-6) |
| 196 | + |
110 | 197 | @patch("sam3.model.io_utils.load_video_frames_from_video_file") |
111 | 198 | def test_unknown_extension_routes_to_video_loader( |
112 | 199 | self, mock_load_video: MagicMock |
|
0 commit comments