File size: 1,750 Bytes
ead274c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """Smoke tests for TransNetV2 wrapper."""
import os
import pytest
from detectors.transnet import detect_transnet
THREESHOT = "/workspace/ComfyUI/input/test_3shots.mp4"
REEL = "/workspace/ComfyUI/input/reel_multishot.mp4"
IMBA = "/workspace/ComfyUI/input/imba.mp4"
@pytest.mark.skipif(not os.path.exists(THREESHOT), reason="synthetic 3-shot missing")
def test_transnet_detects_synthetic_3shot_ground_truth():
boundaries = detect_transnet(THREESHOT, threshold=0.5)
print(f"\n[info] TransNet on test_3shots.mp4: {boundaries}")
assert len(boundaries) == 2, f"expected 2 cuts, got {len(boundaries)}: {boundaries}"
# Ground truth: frame 60 and 120. Tolerance +/- 3 frames.
assert 57 <= boundaries[0] <= 63, f"first cut expected ~60, got {boundaries[0]}"
assert 117 <= boundaries[1] <= 123, f"second cut expected ~120, got {boundaries[1]}"
@pytest.mark.skipif(not os.path.exists(REEL), reason="reel sample missing")
def test_transnet_detects_at_least_one_cut_in_real_reel():
boundaries = detect_transnet(REEL, threshold=0.5)
print(f"\n[info] TransNet on reel_multishot.mp4: {boundaries}")
assert len(boundaries) >= 1, f"expected at least 1 cut in multi-shot reel, got {boundaries}"
# PSD found 1 cut near frame 140. TransNet should agree on at least one cut there or nearby.
@pytest.mark.skipif(not os.path.exists(IMBA), reason="imba sample missing")
def test_transnet_single_shot_returns_empty_or_one():
boundaries = detect_transnet(IMBA, threshold=0.5)
print(f"\n[info] TransNet on imba.mp4: {boundaries}")
# imba.mp4 is a single 8.7s phone video — expect 0, tolerate up to 1 (TransNet is sensitive)
assert len(boundaries) <= 1, f"expected <=1 cut in single-shot video, got {boundaries}"
|