ISR / tests /test_inspection_frames.py
Zhen Ye
feat(inspection): add frame extraction and cropping module
157bd4f
import numpy as np
import pytest
def test_extract_frame_returns_bgr_array(tmp_path):
"""extract_frame should return an HxWx3 BGR numpy array."""
from inspection.frames import extract_frame
# Create a tiny test video (10 frames, 64x48)
import cv2
video_path = str(tmp_path / "test.mp4")
writer = cv2.VideoWriter(
video_path, cv2.VideoWriter_fourcc(*"mp4v"), 30, (64, 48)
)
for i in range(10):
frame = np.full((48, 64, 3), i * 25, dtype=np.uint8)
writer.write(frame)
writer.release()
frame = extract_frame(video_path, 0)
assert isinstance(frame, np.ndarray)
assert frame.shape == (48, 64, 3)
assert frame.dtype == np.uint8
def test_extract_frame_different_indices(tmp_path):
"""Different frame indices should return different pixel data."""
from inspection.frames import extract_frame
import cv2
video_path = str(tmp_path / "test.mp4")
writer = cv2.VideoWriter(
video_path, cv2.VideoWriter_fourcc(*"mp4v"), 30, (64, 48)
)
for i in range(10):
frame = np.full((48, 64, 3), i * 25, dtype=np.uint8)
writer.write(frame)
writer.release()
f0 = extract_frame(video_path, 0)
f5 = extract_frame(video_path, 5)
assert not np.array_equal(f0, f5)
def test_extract_frame_out_of_range(tmp_path):
"""Out-of-range frame index should raise ValueError."""
from inspection.frames import extract_frame
import cv2
video_path = str(tmp_path / "test.mp4")
writer = cv2.VideoWriter(
video_path, cv2.VideoWriter_fourcc(*"mp4v"), 30, (64, 48)
)
for i in range(10):
writer.write(np.zeros((48, 64, 3), dtype=np.uint8))
writer.release()
with pytest.raises(ValueError, match="out of range"):
extract_frame(video_path, 999)
def test_crop_frame_to_bbox():
"""crop_frame should extract the bbox region with padding."""
from inspection.frames import crop_frame
frame = np.zeros((200, 300, 3), dtype=np.uint8)
# Fill a known region with white
frame[50:100, 80:180] = 255
bbox = [80, 50, 180, 100] # x1, y1, x2, y2
crop = crop_frame(frame, bbox, padding=0.0)
assert crop.shape == (50, 100, 3)
assert np.all(crop == 255)
def test_crop_frame_with_padding():
"""Padding should expand the crop region, clamped to frame bounds."""
from inspection.frames import crop_frame
frame = np.zeros((200, 300, 3), dtype=np.uint8)
bbox = [100, 50, 200, 150] # 100x100 box
crop = crop_frame(frame, bbox, padding=0.5)
# 50% padding on a 100x100 box = 50px each side
# Expected: x=[50,250], y=[0,200] (clamped)
assert crop.shape[0] > 100
assert crop.shape[1] > 100
def test_crop_frame_clamped_to_bounds():
"""Padding that exceeds frame bounds should be clamped."""
from inspection.frames import crop_frame
frame = np.zeros((100, 100, 3), dtype=np.uint8)
bbox = [0, 0, 100, 100]
crop = crop_frame(frame, bbox, padding=1.0)
# Should not exceed original frame dimensions
assert crop.shape[0] <= 100
assert crop.shape[1] <= 100