446 lines
18 KiB
Python
446 lines
18 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
Convert phoneme segmented data to phoneme classification dataset
|
||
|
|
将音素分段数据转换为音素分类数据集
|
||
|
|
"""
|
||
|
|
|
||
|
|
import pickle
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from pathlib import Path
|
||
|
|
from collections import defaultdict
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
|
||
|
|
# Add parent directory to path for imports
|
||
|
|
sys.path.append(str(Path(__file__).parent.parent))
|
||
|
|
|
||
|
|
def load_neural_data_for_trial(session, trial_metadata):
|
||
|
|
"""Load neural features for a specific trial"""
|
||
|
|
try:
|
||
|
|
from model_training.evaluate_model_helpers import load_h5py_file
|
||
|
|
import pandas as pd
|
||
|
|
|
||
|
|
# Try to load the session data
|
||
|
|
data_dir = Path(__file__).parent.parent / "data" / "hdf5_data_final"
|
||
|
|
train_file = data_dir / session / "data_train.hdf5"
|
||
|
|
|
||
|
|
if not train_file.exists():
|
||
|
|
return None
|
||
|
|
|
||
|
|
# Load CSV for metadata
|
||
|
|
csv_path = data_dir.parent / "b2txt_dataset_info.csv"
|
||
|
|
if csv_path.exists():
|
||
|
|
b2txt_csv_df = pd.read_csv(csv_path)
|
||
|
|
else:
|
||
|
|
b2txt_csv_df = None
|
||
|
|
|
||
|
|
data = load_h5py_file(str(train_file), b2txt_csv_df)
|
||
|
|
|
||
|
|
# Find the matching trial
|
||
|
|
trial_idx = trial_metadata.get('trial_idx')
|
||
|
|
if trial_idx is not None and trial_idx < len(data['neural_features']):
|
||
|
|
return data['neural_features'][trial_idx]
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Warning: Could not load neural data for {session}, trial {trial_metadata.get('trial_idx', 'unknown')}: {e}")
|
||
|
|
|
||
|
|
return None
|
||
|
|
|
||
|
|
def validate_phoneme_against_ground_truth(segment, ctc_data):
|
||
|
|
"""
|
||
|
|
Validate a phoneme segment against ground truth sequence labels
|
||
|
|
返回: (is_valid, error_reason, ground_truth_phoneme)
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
session = segment['session']
|
||
|
|
trial_idx = segment.get('trial_idx')
|
||
|
|
trial_key = (session, trial_idx)
|
||
|
|
|
||
|
|
if trial_key not in ctc_data:
|
||
|
|
return False, "no_trial_data", None
|
||
|
|
|
||
|
|
trial_data = ctc_data[trial_key]
|
||
|
|
original_sequence = trial_data.get('original_sequence')
|
||
|
|
|
||
|
|
if original_sequence is None:
|
||
|
|
return False, "no_ground_truth", None
|
||
|
|
|
||
|
|
# Convert sequence IDs to phonemes using LOGIT_TO_PHONEME mapping
|
||
|
|
try:
|
||
|
|
from model_training.evaluate_model_helpers import LOGIT_TO_PHONEME
|
||
|
|
except:
|
||
|
|
# Fallback phoneme mapping if import fails
|
||
|
|
LOGIT_TO_PHONEME = {
|
||
|
|
0: 'BLANK', 1: 'AA', 2: 'AE', 3: 'AH', 4: 'AO', 5: 'AW', 6: 'AY',
|
||
|
|
7: 'B', 8: 'CH', 9: 'D', 10: 'DH', 11: 'EH', 12: 'ER', 13: 'EY', 14: 'F',
|
||
|
|
15: 'G', 16: 'HH', 17: 'IH', 18: 'IY', 19: 'JH', 20: 'K', 21: 'L', 22: 'M',
|
||
|
|
23: 'N', 24: 'NG', 25: 'OW', 26: 'OY', 27: 'P', 28: 'R', 29: 'S', 30: 'SH',
|
||
|
|
31: 'T', 32: 'TH', 33: 'UH', 34: 'UW', 35: 'V', 36: 'W', 37: 'Y', 38: 'Z',
|
||
|
|
39: 'ZH', 40: ' | '
|
||
|
|
}
|
||
|
|
|
||
|
|
# Convert ground truth sequence to phonemes (filter out zeros/padding)
|
||
|
|
ground_truth_phonemes = []
|
||
|
|
for seq_id in original_sequence:
|
||
|
|
if seq_id > 0 and seq_id in LOGIT_TO_PHONEME: # Skip padding/blank
|
||
|
|
ground_truth_phonemes.append(LOGIT_TO_PHONEME[seq_id])
|
||
|
|
|
||
|
|
# Find the position of this segment in the predicted sequence
|
||
|
|
predicted_sequence = trial_data.get('predicted_phonemes', [])
|
||
|
|
alignment_info = trial_data.get('alignment_info', [])
|
||
|
|
|
||
|
|
# Find this segment in the alignment info
|
||
|
|
segment_phoneme = segment['phoneme']
|
||
|
|
segment_start = segment['start_time']
|
||
|
|
segment_end = segment['end_time']
|
||
|
|
|
||
|
|
# Find matching alignment segment
|
||
|
|
segment_position = None
|
||
|
|
for i, (phoneme, start, end, conf) in enumerate(alignment_info):
|
||
|
|
if (phoneme == segment_phoneme and
|
||
|
|
start == segment_start and
|
||
|
|
end == segment_end):
|
||
|
|
segment_position = i
|
||
|
|
break
|
||
|
|
|
||
|
|
if segment_position is None:
|
||
|
|
return False, "segment_not_found_in_alignment", None
|
||
|
|
|
||
|
|
# Check if the phoneme at this position matches ground truth
|
||
|
|
if segment_position < len(ground_truth_phonemes):
|
||
|
|
expected_phoneme = ground_truth_phonemes[segment_position]
|
||
|
|
if segment_phoneme == expected_phoneme:
|
||
|
|
return True, "valid", expected_phoneme
|
||
|
|
else:
|
||
|
|
return False, "phoneme_mismatch", expected_phoneme
|
||
|
|
else:
|
||
|
|
return False, "position_out_of_range", None
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
return False, f"validation_error: {str(e)}", None
|
||
|
|
|
||
|
|
def create_phoneme_classification_dataset():
|
||
|
|
"""Create a phoneme classification dataset from segmented data with validation"""
|
||
|
|
|
||
|
|
# Load the latest phoneme dataset
|
||
|
|
data_dir = Path("phoneme_segmented_data")
|
||
|
|
dataset_files = list(data_dir.glob("phoneme_dataset_*.pkl"))
|
||
|
|
|
||
|
|
if not dataset_files:
|
||
|
|
print("No phoneme dataset files found!")
|
||
|
|
return
|
||
|
|
|
||
|
|
latest_dataset = max(dataset_files, key=lambda x: x.stat().st_mtime)
|
||
|
|
print(f"Loading dataset: {latest_dataset.name}")
|
||
|
|
|
||
|
|
with open(latest_dataset, 'rb') as f:
|
||
|
|
phoneme_data = pickle.load(f)
|
||
|
|
|
||
|
|
# Also load the corresponding CTC results for ground truth validation
|
||
|
|
ctc_file = latest_dataset.parent / latest_dataset.name.replace("phoneme_dataset_", "ctc_results_")
|
||
|
|
ctc_data = {}
|
||
|
|
|
||
|
|
if ctc_file.exists():
|
||
|
|
with open(ctc_file, 'rb') as f:
|
||
|
|
ctc_results = pickle.load(f)
|
||
|
|
# Create lookup dictionary for validation
|
||
|
|
for result in ctc_results:
|
||
|
|
key = (result['session'], result['trial_idx'])
|
||
|
|
ctc_data[key] = result
|
||
|
|
|
||
|
|
print(f"Loaded {len(phoneme_data)} phoneme types")
|
||
|
|
print(f"Associated CTC results: {len(ctc_data)} trials")
|
||
|
|
|
||
|
|
# Create classification dataset
|
||
|
|
classification_data = {
|
||
|
|
'features': [], # Neural features for each segment
|
||
|
|
'labels': [], # Phoneme labels
|
||
|
|
'phoneme_to_id': {}, # Phoneme to numeric ID mapping
|
||
|
|
'id_to_phoneme': {}, # Numeric ID to phoneme mapping
|
||
|
|
'metadata': [] # Additional metadata for each sample
|
||
|
|
}
|
||
|
|
|
||
|
|
# Create error tracking
|
||
|
|
error_data = {
|
||
|
|
'incorrect_segments': [], # Incorrect phoneme segments
|
||
|
|
'validation_stats': {} # Validation statistics
|
||
|
|
}
|
||
|
|
|
||
|
|
# Create phoneme to ID mapping
|
||
|
|
unique_phonemes = sorted(phoneme_data.keys())
|
||
|
|
for i, phoneme in enumerate(unique_phonemes):
|
||
|
|
classification_data['phoneme_to_id'][phoneme] = i
|
||
|
|
classification_data['id_to_phoneme'][i] = phoneme
|
||
|
|
|
||
|
|
print(f"\nPhoneme mapping created for {len(unique_phonemes)} phonemes:")
|
||
|
|
for i, phoneme in enumerate(unique_phonemes[:10]): # Show first 10
|
||
|
|
print(f" {i:2d}: '{phoneme}'")
|
||
|
|
if len(unique_phonemes) > 10:
|
||
|
|
print(f" ... and {len(unique_phonemes) - 10} more")
|
||
|
|
|
||
|
|
# Validation and extraction statistics
|
||
|
|
validation_stats = {
|
||
|
|
'total_segments': 0,
|
||
|
|
'valid_segments': 0,
|
||
|
|
'invalid_segments': 0,
|
||
|
|
'discarded_neighbors': 0,
|
||
|
|
'successful_extractions': 0,
|
||
|
|
'error_reasons': defaultdict(int)
|
||
|
|
}
|
||
|
|
|
||
|
|
print(f"\nValidating phoneme segments against ground truth...")
|
||
|
|
|
||
|
|
# First pass: validate all segments and mark invalid ones
|
||
|
|
segment_validity = {} # Maps (phoneme, segment_idx) -> (is_valid, error_reason, ground_truth)
|
||
|
|
|
||
|
|
for phoneme, segments in phoneme_data.items():
|
||
|
|
print(f"Validating '{phoneme}' ({len(segments)} segments)...")
|
||
|
|
|
||
|
|
for segment_idx, segment in enumerate(segments):
|
||
|
|
validation_stats['total_segments'] += 1
|
||
|
|
|
||
|
|
# Validate against ground truth
|
||
|
|
is_valid, error_reason, ground_truth_phoneme = validate_phoneme_against_ground_truth(segment, ctc_data)
|
||
|
|
segment_validity[(phoneme, segment_idx)] = (is_valid, error_reason, ground_truth_phoneme)
|
||
|
|
|
||
|
|
if is_valid:
|
||
|
|
validation_stats['valid_segments'] += 1
|
||
|
|
else:
|
||
|
|
validation_stats['invalid_segments'] += 1
|
||
|
|
validation_stats['error_reasons'][error_reason] += 1
|
||
|
|
|
||
|
|
# Save error information
|
||
|
|
error_data['incorrect_segments'].append({
|
||
|
|
'phoneme': phoneme,
|
||
|
|
'segment_idx': segment_idx,
|
||
|
|
'segment': segment,
|
||
|
|
'predicted_phoneme': phoneme,
|
||
|
|
'ground_truth_phoneme': ground_truth_phoneme,
|
||
|
|
'error_reason': error_reason
|
||
|
|
})
|
||
|
|
|
||
|
|
print(f"\nValidation completed:")
|
||
|
|
print(f" Total segments: {validation_stats['total_segments']}")
|
||
|
|
print(f" Valid segments: {validation_stats['valid_segments']}")
|
||
|
|
print(f" Invalid segments: {validation_stats['invalid_segments']}")
|
||
|
|
print(f" Validation accuracy: {validation_stats['valid_segments']/validation_stats['total_segments']*100:.1f}%")
|
||
|
|
|
||
|
|
print(f"\nError breakdown:")
|
||
|
|
for error_reason, count in validation_stats['error_reasons'].items():
|
||
|
|
print(f" {error_reason}: {count}")
|
||
|
|
|
||
|
|
# Second pass: extract features for valid segments (excluding neighbors of invalid ones)
|
||
|
|
print(f"\nExtracting neural features for validated segments...")
|
||
|
|
|
||
|
|
for phoneme, segments in phoneme_data.items():
|
||
|
|
phoneme_id = classification_data['phoneme_to_id'][phoneme]
|
||
|
|
print(f"Processing '{phoneme}' ({len(segments)} segments)...")
|
||
|
|
|
||
|
|
for segment_idx, segment in enumerate(segments):
|
||
|
|
# Check if this segment is valid
|
||
|
|
is_valid, error_reason, ground_truth_phoneme = segment_validity[(phoneme, segment_idx)]
|
||
|
|
|
||
|
|
# Check if neighboring segments are invalid (discard neighbors)
|
||
|
|
prev_invalid = (segment_idx > 0 and
|
||
|
|
not segment_validity[(phoneme, segment_idx - 1)][0])
|
||
|
|
next_invalid = (segment_idx < len(segments) - 1 and
|
||
|
|
not segment_validity[(phoneme, segment_idx + 1)][0])
|
||
|
|
|
||
|
|
if not is_valid:
|
||
|
|
continue # Skip invalid segments
|
||
|
|
|
||
|
|
if prev_invalid or next_invalid:
|
||
|
|
validation_stats['discarded_neighbors'] += 1
|
||
|
|
continue # Skip neighbors of invalid segments
|
||
|
|
|
||
|
|
# Get trial information
|
||
|
|
session = segment['session']
|
||
|
|
trial_key = (session, segment.get('trial_idx'))
|
||
|
|
|
||
|
|
# Try to get neural data for this trial
|
||
|
|
neural_features = None
|
||
|
|
if trial_key in ctc_data:
|
||
|
|
# We have the trial data, now extract the segment
|
||
|
|
trial_metadata = ctc_data[trial_key]
|
||
|
|
if 'neural_features' in trial_metadata:
|
||
|
|
neural_features = trial_metadata['neural_features']
|
||
|
|
else:
|
||
|
|
# Try to load from HDF5 files
|
||
|
|
neural_features = load_neural_data_for_trial(session, segment)
|
||
|
|
|
||
|
|
if neural_features is not None:
|
||
|
|
# Extract the specific time segment
|
||
|
|
start_time = int(segment['start_time'])
|
||
|
|
end_time = int(segment['end_time'])
|
||
|
|
|
||
|
|
# Ensure valid time range
|
||
|
|
if start_time <= end_time and end_time < len(neural_features):
|
||
|
|
# Extract neural features for this time segment
|
||
|
|
segment_features = neural_features[start_time:end_time+1] # Include end_time
|
||
|
|
|
||
|
|
# Convert to numpy array and handle different cases
|
||
|
|
if isinstance(segment_features, torch.Tensor):
|
||
|
|
segment_features = segment_features.numpy()
|
||
|
|
elif isinstance(segment_features, list):
|
||
|
|
segment_features = np.array(segment_features)
|
||
|
|
|
||
|
|
# For classification, we need a fixed-size feature vector
|
||
|
|
# Option 1: Use mean across time steps
|
||
|
|
if len(segment_features.shape) == 2: # (time, features)
|
||
|
|
feature_vector = np.mean(segment_features, axis=0)
|
||
|
|
elif len(segment_features.shape) == 1: # Already 1D
|
||
|
|
feature_vector = segment_features
|
||
|
|
else:
|
||
|
|
print(f"Unexpected feature shape: {segment_features.shape}")
|
||
|
|
continue
|
||
|
|
|
||
|
|
# Add to dataset
|
||
|
|
classification_data['features'].append(feature_vector)
|
||
|
|
classification_data['labels'].append(phoneme_id)
|
||
|
|
classification_data['metadata'].append({
|
||
|
|
'phoneme': phoneme,
|
||
|
|
'session': session,
|
||
|
|
'trial_num': segment.get('trial_num', -1),
|
||
|
|
'trial_idx': segment.get('trial_idx', -1),
|
||
|
|
'start_time': start_time,
|
||
|
|
'end_time': end_time,
|
||
|
|
'duration': end_time - start_time + 1,
|
||
|
|
'confidence': segment.get('confidence', 0.0),
|
||
|
|
'corpus': segment.get('corpus', 'unknown'),
|
||
|
|
'validated': True
|
||
|
|
})
|
||
|
|
|
||
|
|
validation_stats['successful_extractions'] += 1
|
||
|
|
|
||
|
|
# Progress update
|
||
|
|
if validation_stats['total_segments'] % 1000 == 0:
|
||
|
|
print(f" Processed {validation_stats['total_segments']} segments, extracted {validation_stats['successful_extractions']} features")
|
||
|
|
|
||
|
|
print(f"\nDataset creation completed!")
|
||
|
|
print(f"Total segments processed: {validation_stats['total_segments']}")
|
||
|
|
print(f"Valid segments (excluding neighbors): {validation_stats['valid_segments'] - validation_stats['discarded_neighbors']}")
|
||
|
|
print(f"Discarded neighbor segments: {validation_stats['discarded_neighbors']}")
|
||
|
|
print(f"Successful feature extractions: {validation_stats['successful_extractions']}")
|
||
|
|
print(f"Extraction success rate: {validation_stats['successful_extractions']/(validation_stats['valid_segments']-validation_stats['discarded_neighbors'])*100:.1f}%")
|
||
|
|
|
||
|
|
if validation_stats['successful_extractions'] == 0:
|
||
|
|
print("No features were extracted. Check neural data availability.")
|
||
|
|
return
|
||
|
|
|
||
|
|
# Convert to numpy arrays
|
||
|
|
classification_data['features'] = np.array(classification_data['features'])
|
||
|
|
classification_data['labels'] = np.array(classification_data['labels'])
|
||
|
|
|
||
|
|
print(f"\nFinal validated dataset shape:")
|
||
|
|
print(f"Features: {classification_data['features'].shape}")
|
||
|
|
print(f"Labels: {classification_data['labels'].shape}")
|
||
|
|
|
||
|
|
# Show class distribution
|
||
|
|
print(f"\nClass distribution:")
|
||
|
|
unique_labels, counts = np.unique(classification_data['labels'], return_counts=True)
|
||
|
|
for label_id, count in zip(unique_labels, counts):
|
||
|
|
phoneme = classification_data['id_to_phoneme'][label_id]
|
||
|
|
print(f" {label_id:2d} ('{phoneme}'): {count:4d} samples")
|
||
|
|
|
||
|
|
# Save the classification dataset
|
||
|
|
timestamp = latest_dataset.name.split('_')[-1].replace('.pkl', '')
|
||
|
|
output_file = f"phoneme_classification_dataset_validated_{timestamp}.pkl"
|
||
|
|
output_path = data_dir / output_file
|
||
|
|
|
||
|
|
# Add validation stats to the dataset
|
||
|
|
classification_data['validation_stats'] = validation_stats
|
||
|
|
|
||
|
|
with open(output_path, 'wb') as f:
|
||
|
|
pickle.dump(classification_data, f)
|
||
|
|
|
||
|
|
print(f"\nValidated classification dataset saved to: {output_file}")
|
||
|
|
|
||
|
|
# Save error data separately
|
||
|
|
error_data['validation_stats'] = validation_stats
|
||
|
|
error_file = f"phoneme_validation_errors_{timestamp}.pkl"
|
||
|
|
error_path = data_dir / error_file
|
||
|
|
|
||
|
|
with open(error_path, 'wb') as f:
|
||
|
|
pickle.dump(error_data, f)
|
||
|
|
|
||
|
|
print(f"Validation errors saved to: {error_file}")
|
||
|
|
|
||
|
|
# Create a simple train/test split example
|
||
|
|
create_train_test_split(classification_data, data_dir, timestamp)
|
||
|
|
|
||
|
|
return classification_data
|
||
|
|
|
||
|
|
def create_train_test_split(data, data_dir, timestamp):
|
||
|
|
"""Create train/test split for the classification dataset"""
|
||
|
|
|
||
|
|
from sklearn.model_selection import train_test_split
|
||
|
|
from sklearn.preprocessing import StandardScaler
|
||
|
|
|
||
|
|
print(f"\nCreating train/test split...")
|
||
|
|
|
||
|
|
X = data['features']
|
||
|
|
y = data['labels']
|
||
|
|
metadata = data['metadata']
|
||
|
|
|
||
|
|
# Split by session to avoid data leakage
|
||
|
|
sessions = [meta['session'] for meta in metadata]
|
||
|
|
unique_sessions = list(set(sessions))
|
||
|
|
|
||
|
|
print(f"Available sessions: {len(unique_sessions)}")
|
||
|
|
|
||
|
|
if len(unique_sessions) >= 4:
|
||
|
|
# Use session-based split
|
||
|
|
train_sessions = unique_sessions[:int(len(unique_sessions) * 0.8)]
|
||
|
|
test_sessions = unique_sessions[int(len(unique_sessions) * 0.8):]
|
||
|
|
|
||
|
|
train_indices = [i for i, meta in enumerate(metadata) if meta['session'] in train_sessions]
|
||
|
|
test_indices = [i for i, meta in enumerate(metadata) if meta['session'] in test_sessions]
|
||
|
|
|
||
|
|
X_train, X_test = X[train_indices], X[test_indices]
|
||
|
|
y_train, y_test = y[train_indices], y[test_indices]
|
||
|
|
|
||
|
|
print(f"Session-based split:")
|
||
|
|
print(f" Train sessions: {train_sessions}")
|
||
|
|
print(f" Test sessions: {test_sessions}")
|
||
|
|
else:
|
||
|
|
# Use random split
|
||
|
|
X_train, X_test, y_train, y_test = train_test_split(
|
||
|
|
X, y, test_size=0.2, random_state=42, stratify=y
|
||
|
|
)
|
||
|
|
print(f"Random split (stratified):")
|
||
|
|
|
||
|
|
print(f" Train samples: {len(X_train)}")
|
||
|
|
print(f" Test samples: {len(X_test)}")
|
||
|
|
|
||
|
|
# Standardize features
|
||
|
|
scaler = StandardScaler()
|
||
|
|
X_train_scaled = scaler.fit_transform(X_train)
|
||
|
|
X_test_scaled = scaler.transform(X_test)
|
||
|
|
|
||
|
|
# Save split data
|
||
|
|
split_data = {
|
||
|
|
'X_train': X_train_scaled,
|
||
|
|
'X_test': X_test_scaled,
|
||
|
|
'y_train': y_train,
|
||
|
|
'y_test': y_test,
|
||
|
|
'scaler': scaler,
|
||
|
|
'phoneme_to_id': data['phoneme_to_id'],
|
||
|
|
'id_to_phoneme': data['id_to_phoneme']
|
||
|
|
}
|
||
|
|
|
||
|
|
split_file = f"phoneme_classification_split_{timestamp}.pkl"
|
||
|
|
split_path = data_dir / split_file
|
||
|
|
|
||
|
|
with open(split_path, 'wb') as f:
|
||
|
|
pickle.dump(split_data, f)
|
||
|
|
|
||
|
|
print(f"Train/test split saved to: {split_file}")
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
try:
|
||
|
|
classification_data = create_phoneme_classification_dataset()
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error creating classification dataset: {e}")
|
||
|
|
import traceback
|
||
|
|
traceback.print_exc()
|