Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions sam3/model/io_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand Down
87 changes: 87 additions & 0 deletions test/test_io_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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
Expand Down