diff --git a/sam3/model/io_utils.py b/sam3/model/io_utils.py index 4e76f04f4..2cfd04da2 100644 --- a/sam3/model/io_utils.py +++ b/sam3/model/io_utils.py @@ -342,8 +342,9 @@ def load_video_frames_from_video_file_using_cv2( ) # Convert to tensor - frames_np = np.stack(frames, axis=0).astype(np.float32) # (T, H, W, C) - video_tensor = torch.from_numpy(frames_np).permute(0, 3, 1, 2) # (T, C, H, W) + frames_np = np.stack(frames, axis=0) # (T, H, W, C) + video_tensor = torch.from_numpy(frames_np).permute(0, 3, 1, 2) + video_tensor = video_tensor.to(dtype=torch.float16) # (T, C, H, W) # pyrefly: ignore [bad-assignment] img_mean = torch.tensor(img_mean, dtype=torch.float16).view(1, 3, 1, 1) @@ -357,6 +358,7 @@ def load_video_frames_from_video_file_using_cv2( img_std = img_std.cuda() # normalize by mean and std # pyrefly: ignore [unsupported-operation] + video_tensor /= 255.0 video_tensor -= img_mean # pyrefly: ignore [unsupported-operation] video_tensor /= img_std diff --git a/test/test_io_utils.py b/test/test_io_utils.py index 780b8af78..9fa0c8a40 100644 --- a/test/test_io_utils.py +++ b/test/test_io_utils.py @@ -2,10 +2,14 @@ """Tests for io_utils extensionless video file handling (D99228861).""" +import os import tempfile import unittest from unittest.mock import MagicMock, patch +import numpy as np +import torch + from sam3.model.io_utils import load_video_frames @@ -107,6 +111,89 @@ def test_dummy_video_pattern(self) -> None: self.assertEqual(h, 480) self.assertEqual(w, 640) + def test_cv2_video_file_loader_scales_before_normalization(self) -> None: + """OpenCV video loading should match normalized decoded uint8 frames.""" + try: + import cv2 + except ImportError as exc: + self.skipTest(f"OpenCV is required for this test: {exc}") + + image_size = 8 + source_height = 6 + source_width = 10 + img_mean = (0.5, 0.25, 0.75) + img_std = (0.5, 0.25, 0.25) + yy, xx = np.indices((source_height, source_width), dtype=np.uint16) + frames_rgb = [ + np.stack( + ( + (xx * 23 + yy * 7) % 256, + (xx * 11 + yy * 17 + 3) % 256, + (xx * 5 + yy * 29 + 9) % 256, + ), + axis=-1, + ).astype(np.uint8), + np.stack( + ( + (xx * 13 + yy * 19 + 31) % 256, + (xx * 3 + yy * 41 + 47) % 256, + (xx * 37 + yy * 2 + 61) % 256, + ), + axis=-1, + ).astype(np.uint8), + ] + + with tempfile.TemporaryDirectory() as tmpdir: + video_path = os.path.join(tmpdir, "tiny.avi") + writer = cv2.VideoWriter( + video_path, + cv2.VideoWriter_fourcc(*"MJPG"), + 2.0, + (source_width, source_height), + ) + if not writer.isOpened(): + self.skipTest("OpenCV could not create a temporary MJPG video") + + for frame_rgb in frames_rgb: + writer.write(cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)) + writer.release() + + decoded_frames = [] + cap = cv2.VideoCapture(video_path) + while True: + ret, frame_bgr = cap.read() + if not ret: + break + frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) + decoded_frames.append( + cv2.resize( + frame_rgb, + (image_size, image_size), + interpolation=cv2.INTER_CUBIC, + ) + ) + cap.release() + self.assertEqual(len(decoded_frames), len(frames_rgb)) + + expected = torch.from_numpy(np.stack(decoded_frames, axis=0)) + expected = expected.permute(0, 3, 1, 2).to(dtype=torch.float16) + expected /= 255.0 + expected -= torch.tensor(img_mean, dtype=torch.float16).view(1, 3, 1, 1) + expected /= torch.tensor(img_std, dtype=torch.float16).view(1, 3, 1, 1) + + frames, height, width = load_video_frames( + video_path=video_path, + image_size=image_size, + offload_video_to_cpu=True, + img_mean=img_mean, + img_std=img_std, + video_loader_type="cv2", + ) + + self.assertEqual((height, width), (source_height, source_width)) + self.assertEqual(frames.dtype, torch.float16) + torch.testing.assert_close(frames, expected, rtol=0, atol=1e-6) + @patch("sam3.model.io_utils.load_video_frames_from_video_file") def test_unknown_extension_routes_to_video_loader( self, mock_load_video: MagicMock