ComfyUI-ShotSplitter/tests/test_transnet.py

"""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}"