|
|
|
|
|
""" |
|
|
Main script for ensemble annotation. |
|
|
|
|
|
Usage: |
|
|
python annotate_ensemble.py --input dataset_name --mode balanced --output results.parquet |
|
|
""" |
|
|
|
|
|
import argparse |
|
|
import logging |
|
|
import sys |
|
|
from pathlib import Path |
|
|
import pandas as pd |
|
|
from tqdm import tqdm |
|
|
|
|
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent.parent)) |
|
|
|
|
|
from ensemble_tts.models.emotion import EmotionEnsemble |
|
|
|
|
|
logging.basicConfig( |
|
|
level=logging.INFO, |
|
|
format='%(asctime)s - %(levelname)s - %(message)s' |
|
|
) |
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
def main(): |
|
|
parser = argparse.ArgumentParser(description="Ensemble TTS Annotation") |
|
|
|
|
|
parser.add_argument('--input', type=str, required=True, |
|
|
help='Input dataset name or path') |
|
|
parser.add_argument('--mode', type=str, default='balanced', |
|
|
choices=['quick', 'balanced', 'full'], |
|
|
help='Ensemble mode (quick=2 models, balanced=3, full=5)') |
|
|
parser.add_argument('--output', type=str, default='data/annotated/ensemble_results.parquet', |
|
|
help='Output file path') |
|
|
parser.add_argument('--device', type=str, default='cpu', |
|
|
choices=['cpu', 'cuda'], |
|
|
help='Device to use (cpu or cuda)') |
|
|
parser.add_argument('--voting', type=str, default='weighted', |
|
|
choices=['majority', 'weighted', 'confidence'], |
|
|
help='Voting strategy') |
|
|
parser.add_argument('--max-samples', type=int, default=None, |
|
|
help='Maximum number of samples to process (for testing)') |
|
|
|
|
|
args = parser.parse_args() |
|
|
|
|
|
logger.info("="*60) |
|
|
logger.info("ENSEMBLE TTS ANNOTATION") |
|
|
logger.info("="*60) |
|
|
logger.info(f"Mode: {args.mode}") |
|
|
logger.info(f"Device: {args.device}") |
|
|
logger.info(f"Voting: {args.voting}") |
|
|
|
|
|
|
|
|
logger.info("\n[1/4] Initializing Ensemble...") |
|
|
ensemble = EmotionEnsemble( |
|
|
mode=args.mode, |
|
|
device=args.device, |
|
|
voting_strategy=args.voting |
|
|
) |
|
|
|
|
|
|
|
|
logger.info("\n[2/4] Loading Models...") |
|
|
try: |
|
|
ensemble.load_models() |
|
|
except Exception as e: |
|
|
logger.error(f"Failed to load models: {e}") |
|
|
logger.info("\nPlease ensure all dependencies are installed:") |
|
|
logger.info(" pip install -r requirements.txt") |
|
|
sys.exit(1) |
|
|
|
|
|
|
|
|
logger.info(f"\n[3/4] Loading Dataset: {args.input}") |
|
|
try: |
|
|
from datasets import load_dataset |
|
|
dataset = load_dataset(args.input, split='train') |
|
|
|
|
|
if args.max_samples: |
|
|
dataset = dataset.select(range(min(args.max_samples, len(dataset)))) |
|
|
|
|
|
logger.info(f"Dataset loaded: {len(dataset)} samples") |
|
|
except Exception as e: |
|
|
logger.error(f"Failed to load dataset: {e}") |
|
|
sys.exit(1) |
|
|
|
|
|
|
|
|
logger.info("\n[4/4] Annotating...") |
|
|
results = [] |
|
|
|
|
|
for idx, sample in enumerate(tqdm(dataset, desc="Processing")): |
|
|
try: |
|
|
|
|
|
audio_data = sample.get('audio', {}) |
|
|
if isinstance(audio_data, dict): |
|
|
audio_array = audio_data.get('array', []) |
|
|
sample_rate = audio_data.get('sampling_rate', 16000) |
|
|
else: |
|
|
logger.warning(f"Sample {idx}: No audio data") |
|
|
continue |
|
|
|
|
|
|
|
|
prediction = ensemble.predict(audio_array, sample_rate) |
|
|
|
|
|
|
|
|
result = { |
|
|
'index': idx, |
|
|
'text': sample.get('text', ''), |
|
|
'emotion_label': prediction['label'], |
|
|
'emotion_confidence': prediction['confidence'], |
|
|
'emotion_agreement': prediction['agreement'], |
|
|
'emotion_votes': str(prediction.get('votes', {})), |
|
|
'num_models': len(prediction.get('predictions', [])) |
|
|
} |
|
|
|
|
|
results.append(result) |
|
|
|
|
|
except Exception as e: |
|
|
logger.error(f"Error processing sample {idx}: {e}") |
|
|
continue |
|
|
|
|
|
|
|
|
logger.info(f"\nSaving results to: {args.output}") |
|
|
output_path = Path(args.output) |
|
|
output_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
|
|
|
df = pd.DataFrame(results) |
|
|
df.to_parquet(output_path, index=False) |
|
|
|
|
|
logger.info(f"✅ Saved {len(df)} annotated samples") |
|
|
|
|
|
|
|
|
logger.info("\n" + "="*60) |
|
|
logger.info("STATISTICS") |
|
|
logger.info("="*60) |
|
|
logger.info(f"Total samples: {len(df)}") |
|
|
logger.info(f"Average confidence: {df['emotion_confidence'].mean():.3f}") |
|
|
logger.info(f"Average agreement: {df['emotion_agreement'].mean():.3f}") |
|
|
logger.info("\nEmotion distribution:") |
|
|
logger.info(df['emotion_label'].value_counts()) |
|
|
|
|
|
logger.info("\n✅ DONE!") |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
|
|
|
