Spaces:

BBBAKERY
/

ustwo-api

Sleeping

ustwo-api / scripts /test_english_e2e.py

Deploy from GitHub 2026-04-23T03:56:31Z

c857b85 6 months ago

7.75 kB

#!/usr/bin/env python3
"""영어 E2E 파이프라인 테스트.
RAVDESS WAV + 수동 작성 텍스트로 Stage2 전체 파이프라인을 테스트한다.
감정별 2개씩 14개 세그먼트를 통과시켜 audio + text + fusion 동작을 검증.
Usage:
    python scripts/test_english_e2e.py
"""

from __future__ import annotations

import csv
import json
import logging
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("test_english_e2e")

MANIFEST_PATH = PROJECT_ROOT / "data" / "ravdess" / "manifest.csv"

# 감정별 매칭 텍스트 (RAVDESS는 텍스트 없으므로 수동 작성)
EMOTION_TEXTS = {
    "neutral": [
        "The meeting is scheduled for three o'clock.",
        "I'll pick up the groceries on the way home.",
    ],
    "joy": [
        "I'm so happy to hear from you!",
        "That's wonderful news, I'm thrilled!",
    ],
    "sadness": [
        "I miss you so much, it hurts.",
        "I feel really down today, nothing is going right.",
    ],
    "anger": [
        "I can't believe you did that, I'm so angry!",
        "This is completely unacceptable, stop it now!",
    ],
    "surprise": [
        "Oh my god, I can't believe it!",
        "What?! I never expected that!",
    ],
    "fear": [
        "I'm scared, something doesn't feel right.",
        "Please help, I'm terrified right now.",
    ],
    "disgust": [
        "That's absolutely disgusting, I feel sick.",
        "This is revolting, I can't stand it.",
    ],
}


def pick_samples(manifest_path: Path) -> list[dict]:
    """감정별 2개씩 WAV 선택."""
    by_emotion: dict[str, list[dict]] = {}
    with open(manifest_path) as f:
        for row in csv.DictReader(f):
            emotion = row["emotion"]
            if emotion not in by_emotion:
                by_emotion[emotion] = []
            if len(by_emotion[emotion]) < 2:
                by_emotion[emotion].append(row)

    samples = []
    for emotion in EMOTION_TEXTS:
        if emotion in by_emotion:
            samples.extend(by_emotion[emotion])
    return samples


def run_e2e():
    """E2E 테스트 실행."""
    from src.common.schemas import (
        Models,
        ProcessingInfo,
        Segment,
        Stage1Output,
    )
    from src.stage2.process import process

    if not MANIFEST_PATH.exists():
        logger.error("manifest.csv를 찾을 수 없습니다. 먼저 prepare_ravdess.py를 실행하세요.")
        return False

    samples = pick_samples(MANIFEST_PATH)
    logger.info(f"선택된 테스트 샘플: {len(samples)}개")

    # Stage1Output 생성
    segments = []
    ground_truths = []
    for i, sample in enumerate(samples):
        emotion = sample["emotion"]
        texts = EMOTION_TEXTS[emotion]
        text = texts[i % len(texts)]

        segments.append(Segment(
            segment_id=i,
            speaker_id=f"speaker_{int(sample['actor_id']) % 2}",
            start=float(i * 3.0),
            end=float(i * 3.0 + 3.0),
            text=text,
            language="en",
            audio_path=sample["clean_path"],
            confidence=0.95,
        ))
        ground_truths.append(emotion)

    stage1_output = Stage1Output(
        call_id="ravdess_e2e_test",
        duration=float(len(segments) * 3.0),
        speakers=["speaker_0", "speaker_1"],
        audio_path="data/ravdess/test_call.wav",
        segments=segments,
        processing_info=ProcessingInfo(
            processing_time_sec=1.0,
            models=Models(
                diarization="pyannote/speaker-diarization-3.1",
                asr="large-v3-turbo",
                language_id="whisper",
                alignment="whisperx",
            ),
            device="cpu",
        ),
    )

    # Stage2 config (영어 전용)
    config = {
        "audio_emotion": {
            "model": "iic/emotion2vec_plus_base",
        },
        "text_emotion": {
            "korean_model": "searle-j/kote_for_easygoing_people",
            "english_model": "j-hartmann/emotion-english-distilroberta-base",
        },
        "fusion": {
            "audio_weight": 0.6,
            "text_weight": 0.4,
        },
        "output_path": "data/e2e_test_output.json",
    }

    logger.info("Stage 2 process() 실행 중...")
    output = process(stage1_output, config=config)

    # 검증
    logger.info("\n" + "=" * 70)
    logger.info("E2E 테스트 결과")
    logger.info("=" * 70)

    errors = []

    # 1. EmotionResult 개수 확인
    if len(output.emotions) != len(segments):
        errors.append(f"EmotionResult 개수 불일치: {len(output.emotions)} != {len(segments)}")
    else:
        logger.info(f"[PASS] EmotionResult 개수: {len(output.emotions)}")

    # 2. SpeakerSummary 확인
    if len(output.speaker_summaries) > 0:
        logger.info(f"[PASS] SpeakerSummary 생성: {len(output.speaker_summaries)}명")
    else:
        errors.append("SpeakerSummary가 비어있음")

    # 3. 각 EmotionResult 유효성 확인
    valid_emotions = {"neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"}
    for em in output.emotions:
        if em.fused_emotion not in valid_emotions:
            errors.append(f"유효하지 않은 fused_emotion: {em.fused_emotion}")
        if not (0.0 <= em.fused_confidence <= 1.0):
            errors.append(f"fused_confidence 범위 초과: {em.fused_confidence}")

    if not any("유효하지 않은" in e for e in errors):
        logger.info("[PASS] 모든 EmotionResult 유효")

    # 4. JSON 직렬화 테스트
    try:
        json_str = output.model_dump_json(indent=2)
        roundtrip = json.loads(json_str)
        logger.info(f"[PASS] JSON 직렬화/역직렬화 성공 ({len(json_str)} bytes)")
    except Exception as e:
        errors.append(f"JSON 직렬화 실패: {e}")

    # 5. 결과 테이블
    logger.info(f"\n{'Seg':>3s} | {'Ground Truth':>12s} | {'Audio':>10s} | {'Text':>10s} | {'Fused':>10s} | {'Conf':>5s}")
    logger.info("-" * 65)
    audio_correct = 0
    text_correct = 0
    fused_correct = 0
    for i, (em, gt) in enumerate(zip(output.emotions, ground_truths)):
        a_match = "o" if em.audio_emotion == gt else "x"
        t_match = "o" if em.text_emotion == gt else "x"
        f_match = "o" if em.fused_emotion == gt else "x"
        if em.audio_emotion == gt:
            audio_correct += 1
        if em.text_emotion == gt:
            text_correct += 1
        if em.fused_emotion == gt:
            fused_correct += 1
        logger.info(
            f"{i:3d} | {gt:>12s} | {em.audio_emotion:>8s} {a_match} | "
            f"{em.text_emotion:>8s} {t_match} | {em.fused_emotion:>8s} {f_match} | "
            f"{em.fused_confidence:.3f}"
        )

    n = len(ground_truths)
    logger.info(f"\nAccuracy: audio={audio_correct}/{n} text={text_correct}/{n} fused={fused_correct}/{n}")

    # SpeakerSummary 출력
    logger.info("\nSpeaker Summaries:")
    for spk, summary in output.speaker_summaries.items():
        logger.info(f"  {spk}: dominant={summary.dominant_emotion}, "
                     f"conf={summary.avg_confidence:.3f}, "
                     f"dist={summary.emotion_distribution}")

    if errors:
        logger.error(f"\nFAILED — {len(errors)} errors:")
        for e in errors:
            logger.error(f"  - {e}")
        return False

    logger.info(f"\n{'='*70}")
    logger.info("E2E 테스트 PASS — 영어 파이프라인 정상 동작")
    logger.info(f"{'='*70}")
    return True


if __name__ == "__main__":
    success = run_e2e()
    sys.exit(0 if success else 1)

	#!/usr/bin/env python3
	"""영어 E2E 파이프라인 테스트.

	RAVDESS WAV + 수동 작성 텍스트로 Stage2 전체 파이프라인을 테스트한다.
	감정별 2개씩 14개 세그먼트를 통과시켜 audio + text + fusion 동작을 검증.

	Usage:
	python scripts/test_english_e2e.py
	"""

	from __future__ import annotations

	import csv
	import json
	import logging
	import sys
	from pathlib import Path

	PROJECT_ROOT = Path(__file__).parent.parent
	sys.path.insert(0, str(PROJECT_ROOT))

	logging.basicConfig(
	level=logging.INFO,
	format="%(asctime)s - %(levelname)s - %(message)s",
	)
	logger = logging.getLogger("test_english_e2e")

	MANIFEST_PATH = PROJECT_ROOT / "data" / "ravdess" / "manifest.csv"

	# 감정별 매칭 텍스트 (RAVDESS는 텍스트 없으므로 수동 작성)
	EMOTION_TEXTS = {
	"neutral": [
	"The meeting is scheduled for three o'clock.",
	"I'll pick up the groceries on the way home.",
	],
	"joy": [
	"I'm so happy to hear from you!",
	"That's wonderful news, I'm thrilled!",
	],
	"sadness": [
	"I miss you so much, it hurts.",
	"I feel really down today, nothing is going right.",
	],
	"anger": [
	"I can't believe you did that, I'm so angry!",
	"This is completely unacceptable, stop it now!",
	],
	"surprise": [
	"Oh my god, I can't believe it!",
	"What?! I never expected that!",
	],
	"fear": [
	"I'm scared, something doesn't feel right.",
	"Please help, I'm terrified right now.",
	],
	"disgust": [
	"That's absolutely disgusting, I feel sick.",
	"This is revolting, I can't stand it.",
	],
	}


	def pick_samples(manifest_path: Path) -> list[dict]:
	"""감정별 2개씩 WAV 선택."""
	by_emotion: dict[str, list[dict]] = {}
	with open(manifest_path) as f:
	for row in csv.DictReader(f):
	emotion = row["emotion"]
	if emotion not in by_emotion:
	by_emotion[emotion] = []
	if len(by_emotion[emotion]) < 2:
	by_emotion[emotion].append(row)

	samples = []
	for emotion in EMOTION_TEXTS:
	if emotion in by_emotion:
	samples.extend(by_emotion[emotion])
	return samples


	def run_e2e():
	"""E2E 테스트 실행."""
	from src.common.schemas import (
	Models,
	ProcessingInfo,
	Segment,
	Stage1Output,
	)
	from src.stage2.process import process

	if not MANIFEST_PATH.exists():
	logger.error("manifest.csv를 찾을 수 없습니다. 먼저 prepare_ravdess.py를 실행하세요.")
	return False

	samples = pick_samples(MANIFEST_PATH)
	logger.info(f"선택된 테스트 샘플: {len(samples)}개")

	# Stage1Output 생성
	segments = []
	ground_truths = []
	for i, sample in enumerate(samples):
	emotion = sample["emotion"]
	texts = EMOTION_TEXTS[emotion]
	text = texts[i % len(texts)]

	segments.append(Segment(
	segment_id=i,
	speaker_id=f"speaker_{int(sample['actor_id']) % 2}",
	start=float(i * 3.0),
	end=float(i * 3.0 + 3.0),
	text=text,
	language="en",
	audio_path=sample["clean_path"],
	confidence=0.95,
	))
	ground_truths.append(emotion)

	stage1_output = Stage1Output(
	call_id="ravdess_e2e_test",
	duration=float(len(segments) * 3.0),
	speakers=["speaker_0", "speaker_1"],
	audio_path="data/ravdess/test_call.wav",
	segments=segments,
	processing_info=ProcessingInfo(
	processing_time_sec=1.0,
	models=Models(
	diarization="pyannote/speaker-diarization-3.1",
	asr="large-v3-turbo",
	language_id="whisper",
	alignment="whisperx",
	),
	device="cpu",
	),
	)

	# Stage2 config (영어 전용)
	config = {
	"audio_emotion": {
	"model": "iic/emotion2vec_plus_base",
	},
	"text_emotion": {
	"korean_model": "searle-j/kote_for_easygoing_people",
	"english_model": "j-hartmann/emotion-english-distilroberta-base",
	},
	"fusion": {
	"audio_weight": 0.6,
	"text_weight": 0.4,
	},
	"output_path": "data/e2e_test_output.json",
	}

	logger.info("Stage 2 process() 실행 중...")
	output = process(stage1_output, config=config)

	# 검증
	logger.info("\n" + "=" * 70)
	logger.info("E2E 테스트 결과")
	logger.info("=" * 70)

	errors = []

	# 1. EmotionResult 개수 확인
	if len(output.emotions) != len(segments):
	errors.append(f"EmotionResult 개수 불일치: {len(output.emotions)} != {len(segments)}")
	else:
	logger.info(f"[PASS] EmotionResult 개수: {len(output.emotions)}")

	# 2. SpeakerSummary 확인
	if len(output.speaker_summaries) > 0:
	logger.info(f"[PASS] SpeakerSummary 생성: {len(output.speaker_summaries)}명")
	else:
	errors.append("SpeakerSummary가 비어있음")

	# 3. 각 EmotionResult 유효성 확인
	valid_emotions = {"neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"}
	for em in output.emotions:
	if em.fused_emotion not in valid_emotions:
	errors.append(f"유효하지 않은 fused_emotion: {em.fused_emotion}")
	if not (0.0 <= em.fused_confidence <= 1.0):
	errors.append(f"fused_confidence 범위 초과: {em.fused_confidence}")

	if not any("유효하지 않은" in e for e in errors):
	logger.info("[PASS] 모든 EmotionResult 유효")

	# 4. JSON 직렬화 테스트
	try:
	json_str = output.model_dump_json(indent=2)
	roundtrip = json.loads(json_str)
	logger.info(f"[PASS] JSON 직렬화/역직렬화 성공 ({len(json_str)} bytes)")
	except Exception as e:
	errors.append(f"JSON 직렬화 실패: {e}")

	# 5. 결과 테이블
	logger.info(f"\n{'Seg':>3s} \| {'Ground Truth':>12s} \| {'Audio':>10s} \| {'Text':>10s} \| {'Fused':>10s} \| {'Conf':>5s}")
	logger.info("-" * 65)
	audio_correct = 0
	text_correct = 0
	fused_correct = 0
	for i, (em, gt) in enumerate(zip(output.emotions, ground_truths)):
	a_match = "o" if em.audio_emotion == gt else "x"
	t_match = "o" if em.text_emotion == gt else "x"
	f_match = "o" if em.fused_emotion == gt else "x"
	if em.audio_emotion == gt:
	audio_correct += 1
	if em.text_emotion == gt:
	text_correct += 1
	if em.fused_emotion == gt:
	fused_correct += 1
	logger.info(
	f"{i:3d} \| {gt:>12s} \| {em.audio_emotion:>8s} {a_match} \| "
	f"{em.text_emotion:>8s} {t_match} \| {em.fused_emotion:>8s} {f_match} \| "
	f"{em.fused_confidence:.3f}"
	)

	n = len(ground_truths)
	logger.info(f"\nAccuracy: audio={audio_correct}/{n} text={text_correct}/{n} fused={fused_correct}/{n}")

	# SpeakerSummary 출력
	logger.info("\nSpeaker Summaries:")
	for spk, summary in output.speaker_summaries.items():
	logger.info(f" {spk}: dominant={summary.dominant_emotion}, "
	f"conf={summary.avg_confidence:.3f}, "
	f"dist={summary.emotion_distribution}")

	if errors:
	logger.error(f"\nFAILED — {len(errors)} errors:")
	for e in errors:
	logger.error(f" - {e}")
	return False

	logger.info(f"\n{'='*70}")
	logger.info("E2E 테스트 PASS — 영어 파이프라인 정상 동작")
	logger.info(f"{'='*70}")
	return True


	if __name__ == "__main__":
	success = run_e2e()
	sys.exit(0 if success else 1)