livekit-cameras/tests/test_speaker_chrome.py

"""Green speaking chrome on the KMS grid (LiveKit active-speaker)."""
import numpy as np

from display_grid import FRAME_RAISE, FRAME_SPEAK, CameraGrid, bgra_to_i420, bgra_to_yuv_pixel


def test_set_speakers_paints_green_i420_border():
    g = CameraGrid(width=1920, height=1080, cols=5, rows=4)
    src = np.zeros((360, 640, 4), dtype=np.uint8)
    src[:] = (40, 40, 40, 255)
    g.set_frame_i420("cam-07", bgra_to_i420(src), 640, 360)
    before = np.frombuffer(g.render_i420(), dtype=np.uint8).copy()
    g.set_speakers(["cam-07"])
    assert g.speaking("cam-07") is True
    after = np.frombuffer(g.render_i420(), dtype=np.uint8)
    assert not np.array_equal(before, after)
    idx = g.identities.index("cam-07")
    ox, oy = g.tile_origin(idx)
    yplane = after[: g.width * g.height].reshape(g.height, g.width)
    yv, _, _ = bgra_to_yuv_pixel(FRAME_SPEAK)
    assert int(yplane[oy, ox + g.tile_w // 2]) == yv
    # Extra thickness lives in the gutter, not on the feed.
    assert int(yplane[oy - 2, ox + g.tile_w // 2]) == yv
    ix, iy, iw, _ih = g.inner_rect(idx)
    assert int(yplane[iy, ix + iw // 2]) != yv


def test_speaking_overrides_yellow_then_yellow_returns():
    g = CameraGrid(width=1920, height=1080, cols=5, rows=4)
    src = np.zeros((360, 640, 4), dtype=np.uint8)
    g.set_frame_i420("cam-07", bgra_to_i420(src), 640, 360)
    g.set_highlight("cam-07", True)
    g.set_speakers(["cam-07"])
    idx = g.identities.index("cam-07")
    ox, oy = g.tile_origin(idx)
    yplane = np.frombuffer(g.render_i420(), dtype=np.uint8)[: g.width * g.height].reshape(
        g.height, g.width)
    y_speak, _, _ = bgra_to_yuv_pixel(FRAME_SPEAK)
    y_raise, _, _ = bgra_to_yuv_pixel(FRAME_RAISE)
    assert int(yplane[oy, ox + g.tile_w // 2]) == y_speak
    g.set_speakers([])
    yplane = np.frombuffer(g.render_i420(), dtype=np.uint8)[: g.width * g.height].reshape(
        g.height, g.width)
    assert g.speaking("cam-07") is False
    assert g.highlighted("cam-07") is True
    assert int(yplane[oy, ox + g.tile_w // 2]) == y_raise


def test_clear_drops_speaking():
    g = CameraGrid(width=1920, height=1080, cols=5, rows=4)
    src = np.zeros((360, 640, 4), dtype=np.uint8)
    g.set_frame_i420("cam-01", bgra_to_i420(src), 640, 360)
    g.set_speakers(["cam-01"])
    g.clear("cam-01")
    assert g.speaking("cam-01") is False


class _P:
    def __init__(self, identity, keep=True):
        self.identity = identity
        self.keep = keep


def test_active_speaker_ids_keeps_order_and_filter():
    from display_grid import active_speaker_ids, primary_grid_speaker
    speakers = [_P("cam-07"), _P("rally", False), _P("cam-03"), _P("")]
    assert active_speaker_ids(speakers, want=lambda p: p.keep) == ["cam-07", "cam-03"]
    ids = ["rally", "cam-03", "cam-07"]
    assert primary_grid_speaker(ids, ["cam-01", "cam-03", "cam-07"]) == ["cam-03"]
    assert primary_grid_speaker(["rally"], ["cam-01"]) == []


def test_select_grid_speakers_drops_floor_noise():
    from display_grid import select_grid_speakers
    assert select_grid_speakers(
        [("cam-01", -34.0), ("cam-02", -36.0)], min_db=-24.0, margin_db=6.0,
    ) == []


def test_select_grid_speakers_bleed_keeps_louder_only():
    from display_grid import select_grid_speakers
    assert select_grid_speakers(
        [("cam-07", -12.0), ("cam-08", -22.0)], min_db=-24.0, margin_db=6.0,
    ) == ["cam-07"]


def test_select_grid_speakers_two_talkers_both_kept():
    from display_grid import select_grid_speakers
    assert select_grid_speakers(
        [("cam-07", -12.0), ("cam-08", -14.0)], min_db=-24.0, margin_db=6.0,
    ) == ["cam-07", "cam-08"]


def test_select_grid_speakers_normal_speech_above_mic_floor():
    """Conversational C920 speech is often ~-24 dBFS; absolute -20 missed it."""
    from display_grid import select_grid_speakers
    floors = {"cam-07": -34.0, "cam-08": -34.0}
    assert select_grid_speakers(
        [("cam-07", -24.0), ("cam-08", -33.0)],
        min_db=-40.0, margin_db=6.0, floors=floors, rise_db=6.0,
    ) == ["cam-07"]


def test_select_grid_speakers_two_talkers_each_above_own_floor():
    from display_grid import select_grid_speakers
    floors = {"cam-07": -34.0, "cam-08": -36.0}
    assert select_grid_speakers(
        [("cam-07", -24.0), ("cam-08", -26.0)],
        min_db=-40.0, margin_db=6.0, floors=floors, rise_db=6.0,
    ) == ["cam-07", "cam-08"]


def test_select_quiet_mic_speech_not_cut_by_abs_min():
    """Quiet C920 idle is ~-45; +3 dB speech is still below -40 dBFS."""
    from display_grid import select_grid_speakers
    floors = {"cam-10": -45.0, "cam-01": -37.0}
    assert select_grid_speakers(
        [("cam-10", -41.5), ("cam-01", -36.5)],
        min_db=-40.0, margin_db=6.0, floors=floors, rise_db=3.0,
    ) == ["cam-10"]


def test_update_noise_floor_tracks_idle_not_speech():
    from display_grid import update_noise_floor
    fl = update_noise_floor(-34.0, -33.5, speaking=False, rise_db=4.0)
    assert -34.0 < fl < -33.5
    stuck = update_noise_floor(-34.0, -20.0, speaking=True, rise_db=4.0)
    assert stuck == -34.0


def test_confirm_speakers_needs_two_hops_new_talker():
    from display_grid import confirm_speakers
    first, counts = confirm_speakers({}, ["cam-07"], need=2)
    assert first == []
    second, counts = confirm_speakers(counts, ["cam-07"], need=2)
    assert second == ["cam-07"]
    # already on: one hop of continued speech keeps chrome
    stay, _ = confirm_speakers(counts, ["cam-07"], need=2, already=["cam-07"])
    assert stay == ["cam-07"]
    # a one-hop idle spike on another seat does not light
    spike, _ = confirm_speakers({}, ["cam-01"], need=2, already=["cam-07"])
    assert spike == []


def test_waveform_similarity_high_for_delayed_copy():
    from display_grid import waveform_similarity
    t = np.linspace(0, 1, 1600, endpoint=False)
    a = np.sin(2 * np.pi * 220 * t).astype(np.float32)
    b = np.roll(a * 0.4, 80)
    assert waveform_similarity(a, b) > 0.85


def test_waveform_similarity_low_for_independent_signals():
    from display_grid import waveform_similarity
    rng = np.random.default_rng(0)
    a = rng.normal(size=1600).astype(np.float32)
    b = rng.normal(size=1600).astype(np.float32)
    assert waveform_similarity(a, b) < 0.35


def test_collapse_bleed_drops_similar_quieter_mic():
    from display_grid import collapse_bleed
    t = np.linspace(0, 1, 1600, endpoint=False)
    src = np.sin(2 * np.pi * 180 * t).astype(np.float32)
    waves = {"cam-07": src, "cam-08": np.roll(src * 0.3, 40)}
    levels = {"cam-07": -12.0, "cam-08": -18.0}
    assert collapse_bleed(["cam-07", "cam-08"], waves, levels, corr_min=0.4) == ["cam-07"]


def test_collapse_bleed_keeps_two_unlike_talkers():
    from display_grid import collapse_bleed
    t = np.linspace(0, 1, 1600, endpoint=False)
    a = np.sin(2 * np.pi * 180 * t).astype(np.float32)
    b = np.sin(2 * np.pi * 440 * t + 1.3).astype(np.float32)
    waves = {"cam-07": a, "cam-08": b}
    levels = {"cam-07": -12.0, "cam-08": -13.0}
    assert collapse_bleed(["cam-07", "cam-08"], waves, levels, corr_min=0.6) == [
        "cam-07", "cam-08",
    ]


def test_hold_speakers_adds_second_talker_immediately():
    from display_grid import hold_speakers
    cur, until = hold_speakers(
        ["cam-07"], ["cam-07", "cam-08"], prev_peak=-12.0, new_peak=-12.0,
        now=10.0, hold_until=10.5, hold_s=0.6, switch_db=3.0)
    assert cur == ["cam-07", "cam-08"]


def test_hold_speakers_ignores_similar_bleed_flicker():
    from display_grid import hold_speakers
    cur, until = hold_speakers(
        ["cam-07"], ["cam-08"], prev_peak=-12.0, new_peak=-12.5,
        now=10.2, hold_until=10.8, hold_s=0.6, switch_db=3.0)
    assert cur == ["cam-07"]
    cur, _ = hold_speakers(
        ["cam-07"], ["cam-08"], prev_peak=-12.0, new_peak=-8.0,
        now=10.2, hold_until=10.8, hold_s=0.6, switch_db=3.0)
    assert cur == ["cam-08"]