328 lines
13 KiB
Python
328 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
将PKL文件中的输出时间戳转换为原始数据时间戳
|
||
支持三种映射方式：简单映射、保守映射、可能映射
|
||
"""
|
||
|
||
import pickle
|
||
import numpy as np
|
||
from pathlib import Path
|
||
from collections import defaultdict
|
||
import argparse
|
||
from datetime import datetime
|
||
|
||
# 滑动窗口参数 (来自 rnn_args.yaml)
|
||
PATCH_SIZE = 14 # 滑动窗口大小
|
||
PATCH_STRIDE = 4 # 滑动窗口步长
|
||
ORIGINAL_BIN_MS = 20 # 原始时间bin大小(ms)
|
||
|
||
def convert_output_timestamp_to_original(output_start, output_end, patch_size=PATCH_SIZE, patch_stride=PATCH_STRIDE):
|
||
"""
|
||
将输出时间戳转换为原始数据时间戳，提供三种映射方式
|
||
|
||
Args:
|
||
output_start: 输出序列中的开始时间步
|
||
output_end: 输出序列中的结束时间步
|
||
patch_size: 滑动窗口大小
|
||
patch_stride: 滑动窗口步长
|
||
|
||
Returns:
|
||
dict: 包含三种映射方式的原始时间戳信息
|
||
"""
|
||
|
||
# 计算输出时间步对应的原始时间步中心位置
|
||
original_start_center = output_start * patch_stride + (patch_size - 1) / 2
|
||
original_end_center = output_end * patch_stride + (patch_size - 1) / 2
|
||
|
||
# 1. 简单映射：使用中心位置
|
||
simple_start = int(round(original_start_center))
|
||
simple_end = int(round(original_end_center))
|
||
|
||
# 2. 保守映射：考虑完整的patch范围
|
||
patch_start_first = output_start * patch_stride
|
||
patch_end_first = patch_start_first + patch_size - 1
|
||
patch_start_last = output_end * patch_stride
|
||
patch_end_last = patch_start_last + patch_size - 1
|
||
|
||
conservative_start = patch_start_first
|
||
conservative_end = patch_end_last
|
||
|
||
# 3. 可能映射：基于中心位置但考虑patch边界的调整范围
|
||
likely_start = max(patch_start_first, int(original_start_center - patch_stride/2))
|
||
likely_end = min(patch_end_last, int(original_end_center + patch_stride/2))
|
||
|
||
return {
|
||
'simple_mapping': (simple_start, simple_end),
|
||
'conservative_mapping': (conservative_start, conservative_end),
|
||
'likely_mapping': (likely_start, likely_end),
|
||
'metadata': {
|
||
'output_range': (output_start, output_end),
|
||
'output_duration': output_end - output_start + 1,
|
||
'center_positions': (original_start_center, original_end_center),
|
||
'patch_ranges': {
|
||
'first_patch': (patch_start_first, patch_end_first),
|
||
'last_patch': (patch_start_last, patch_end_last)
|
||
}
|
||
}
|
||
}
|
||
|
||
def process_phoneme_dataset(input_file_path, output_file_path):
|
||
"""
|
||
处理phoneme dataset，转换时间戳并保存结果
|
||
|
||
Args:
|
||
input_file_path: 输入PKL文件路径
|
||
output_file_path: 输出PKL文件路径
|
||
"""
|
||
|
||
print(f"=== 处理phoneme dataset ===")
|
||
print(f"输入文件: {input_file_path}")
|
||
print(f"输出文件: {output_file_path}")
|
||
|
||
# 加载原始数据
|
||
with open(input_file_path, 'rb') as f:
|
||
phoneme_data = pickle.load(f)
|
||
|
||
print(f"原始数据类型: {type(phoneme_data)}")
|
||
if isinstance(phoneme_data, dict):
|
||
print(f"音素数量: {len(phoneme_data)}")
|
||
total_segments = sum(len(segments) for segments in phoneme_data.values())
|
||
print(f"总segment数量: {total_segments}")
|
||
|
||
# 转换数据结构
|
||
converted_data = {}
|
||
conversion_stats = {
|
||
'total_segments': 0,
|
||
'conversion_errors': 0,
|
||
'phoneme_counts': defaultdict(int)
|
||
}
|
||
|
||
for phoneme, segments in phoneme_data.items():
|
||
converted_segments = []
|
||
|
||
for segment in segments:
|
||
try:
|
||
# 获取原始时间戳
|
||
output_start = segment['start_time']
|
||
output_end = segment['end_time']
|
||
|
||
# 转换时间戳
|
||
conversion = convert_output_timestamp_to_original(output_start, output_end)
|
||
|
||
# 创建新的segment，保留原有信息并添加转换结果
|
||
new_segment = segment.copy()
|
||
|
||
# 添加原始时间戳字段
|
||
new_segment['original_timestamps'] = {
|
||
'simple_start': conversion['simple_mapping'][0],
|
||
'simple_end': conversion['simple_mapping'][1],
|
||
'conservative_start': conversion['conservative_mapping'][0],
|
||
'conservative_end': conversion['conservative_mapping'][1],
|
||
'likely_start': conversion['likely_mapping'][0],
|
||
'likely_end': conversion['likely_mapping'][1]
|
||
}
|
||
|
||
# 添加时长信息（毫秒）
|
||
output_duration_ms = conversion['metadata']['output_duration'] * PATCH_STRIDE * ORIGINAL_BIN_MS
|
||
new_segment['duration_info'] = {
|
||
'output_duration_steps': conversion['metadata']['output_duration'],
|
||
'output_duration_ms': output_duration_ms,
|
||
'simple_duration_steps': conversion['simple_mapping'][1] - conversion['simple_mapping'][0] + 1,
|
||
'conservative_duration_steps': conversion['conservative_mapping'][1] - conversion['conservative_mapping'][0] + 1,
|
||
'likely_duration_steps': conversion['likely_mapping'][1] - conversion['likely_mapping'][0] + 1
|
||
}
|
||
|
||
# 添加转换元数据
|
||
new_segment['conversion_metadata'] = conversion['metadata']
|
||
|
||
converted_segments.append(new_segment)
|
||
conversion_stats['total_segments'] += 1
|
||
conversion_stats['phoneme_counts'][phoneme] += 1
|
||
|
||
except Exception as e:
|
||
print(f"转换segment时出错 (phoneme: {phoneme}): {e}")
|
||
conversion_stats['conversion_errors'] += 1
|
||
continue
|
||
|
||
converted_data[phoneme] = converted_segments
|
||
|
||
# 添加转换参数信息
|
||
conversion_info = {
|
||
'conversion_timestamp': datetime.now().isoformat(),
|
||
'parameters': {
|
||
'patch_size': PATCH_SIZE,
|
||
'patch_stride': PATCH_STRIDE,
|
||
'original_bin_ms': ORIGINAL_BIN_MS
|
||
},
|
||
'statistics': dict(conversion_stats),
|
||
'mapping_methods': {
|
||
'simple_mapping': '基于输出时间步中心位置的映射',
|
||
'conservative_mapping': '基于完整patch范围的保守映射',
|
||
'likely_mapping': '基于中心位置但考虑patch边界的调整映射'
|
||
}
|
||
}
|
||
|
||
# 保存结果
|
||
output_data = {
|
||
'phoneme_data': converted_data,
|
||
'conversion_info': conversion_info
|
||
}
|
||
|
||
with open(output_file_path, 'wb') as f:
|
||
pickle.dump(output_data, f)
|
||
|
||
print(f"\n=== 转换完成 ===")
|
||
print(f"成功转换: {conversion_stats['total_segments']} 个segments")
|
||
print(f"转换错误: {conversion_stats['conversion_errors']} 个segments")
|
||
print(f"音素分布:")
|
||
|
||
sorted_phonemes = sorted(conversion_stats['phoneme_counts'].items(),
|
||
key=lambda x: x[1], reverse=True)
|
||
for phoneme, count in sorted_phonemes[:10]:
|
||
print(f" {phoneme}: {count} segments")
|
||
|
||
return output_data
|
||
|
||
def analyze_conversion_results(converted_data_path):
|
||
"""
|
||
分析转换结果的统计信息
|
||
|
||
Args:
|
||
converted_data_path: 转换后的PKL文件路径
|
||
"""
|
||
|
||
print(f"\n=== 分析转换结果 ===")
|
||
|
||
with open(converted_data_path, 'rb') as f:
|
||
data = pickle.load(f)
|
||
|
||
phoneme_data = data['phoneme_data']
|
||
conversion_info = data['conversion_info']
|
||
|
||
print(f"转换时间: {conversion_info['conversion_timestamp']}")
|
||
print(f"转换参数: {conversion_info['parameters']}")
|
||
|
||
# 统计分析
|
||
all_segments = []
|
||
for phoneme, segments in phoneme_data.items():
|
||
all_segments.extend(segments)
|
||
|
||
print(f"总segment数量: {len(all_segments)}")
|
||
|
||
# 分析时长分布
|
||
output_durations = []
|
||
simple_durations = []
|
||
conservative_durations = []
|
||
likely_durations = []
|
||
|
||
for segment in all_segments:
|
||
duration_info = segment['duration_info']
|
||
output_durations.append(duration_info['output_duration_steps'])
|
||
simple_durations.append(duration_info['simple_duration_steps'])
|
||
conservative_durations.append(duration_info['conservative_duration_steps'])
|
||
likely_durations.append(duration_info['likely_duration_steps'])
|
||
|
||
print(f"\n时长统计 (时间步):")
|
||
print(f"输出时长: 平均 {np.mean(output_durations):.1f}, 中位数 {np.median(output_durations):.1f}")
|
||
print(f"简单映射时长: 平均 {np.mean(simple_durations):.1f}, 中位数 {np.median(simple_durations):.1f}")
|
||
print(f"保守映射时长: 平均 {np.mean(conservative_durations):.1f}, 中位数 {np.median(conservative_durations):.1f}")
|
||
print(f"可能映射时长: 平均 {np.mean(likely_durations):.1f}, 中位数 {np.median(likely_durations):.1f}")
|
||
|
||
# 计算映射比例
|
||
simple_ratios = [s/o for s, o in zip(simple_durations, output_durations) if o > 0]
|
||
conservative_ratios = [c/o for c, o in zip(conservative_durations, output_durations) if o > 0]
|
||
likely_ratios = [l/o for l, o in zip(likely_durations, output_durations) if o > 0]
|
||
|
||
print(f"\n映射比例 (原始/输出):")
|
||
print(f"简单映射: 平均 {np.mean(simple_ratios):.1f}x")
|
||
print(f"保守映射: 平均 {np.mean(conservative_ratios):.1f}x")
|
||
print(f"可能映射: 平均 {np.mean(likely_ratios):.1f}x")
|
||
|
||
# 显示几个示例
|
||
print(f"\n=== 转换示例 ===")
|
||
sample_segments = all_segments[:5]
|
||
|
||
print(f"{'音素':4s} {'输出时间戳':12s} {'简单映射':12s} {'保守映射':12s} {'可能映射':12s} {'时长(ms)':8s}")
|
||
print("-" * 70)
|
||
|
||
for segment in sample_segments:
|
||
phoneme = segment['phoneme']
|
||
output_range = f"{segment['start_time']}-{segment['end_time']}"
|
||
|
||
timestamps = segment['original_timestamps']
|
||
simple = f"{timestamps['simple_start']}-{timestamps['simple_end']}"
|
||
conservative = f"{timestamps['conservative_start']}-{timestamps['conservative_end']}"
|
||
likely = f"{timestamps['likely_start']}-{timestamps['likely_end']}"
|
||
|
||
duration_ms = segment['duration_info']['output_duration_ms']
|
||
|
||
print(f"{phoneme:4s} {output_range:12s} {simple:12s} {conservative:12s} {likely:12s} {duration_ms:6.0f}")
|
||
|
||
def main():
|
||
"""主函数"""
|
||
parser = argparse.ArgumentParser(description='转换phoneme dataset的时间戳')
|
||
parser.add_argument('--input_dir', type=str, default='../phoneme_segmented_data',
|
||
help='输入PKL文件目录')
|
||
parser.add_argument('--output_dir', type=str, default='../phoneme_segmented_data',
|
||
help='输出PKL文件目录')
|
||
parser.add_argument('--input_file', type=str, default=None,
|
||
help='指定输入文件名，如果不指定则处理最新的phoneme_dataset文件')
|
||
parser.add_argument('--output_suffix', type=str, default='_with_original_timestamps',
|
||
help='输出文件后缀')
|
||
parser.add_argument('--analyze_only', action='store_true',
|
||
help='只分析已存在的转换结果，不进行转换')
|
||
|
||
args = parser.parse_args()
|
||
|
||
input_dir = Path(args.input_dir)
|
||
output_dir = Path(args.output_dir)
|
||
|
||
# 确保输出目录存在
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
if args.analyze_only:
|
||
# 只进行分析
|
||
converted_files = list(output_dir.glob(f"*{args.output_suffix}.pkl"))
|
||
if converted_files:
|
||
latest_converted = max(converted_files, key=lambda x: x.stat().st_mtime)
|
||
analyze_conversion_results(latest_converted)
|
||
else:
|
||
print("未找到转换后的文件")
|
||
return
|
||
|
||
# 找到输入文件
|
||
if args.input_file:
|
||
input_file = input_dir / args.input_file
|
||
else:
|
||
# 找到最新的phoneme_dataset文件
|
||
phoneme_files = list(input_dir.glob("phoneme_dataset_*.pkl"))
|
||
if not phoneme_files:
|
||
print(f"在目录 {input_dir} 中未找到phoneme_dataset文件")
|
||
return
|
||
|
||
input_file = max(phoneme_files, key=lambda x: x.stat().st_mtime)
|
||
|
||
if not input_file.exists():
|
||
print(f"输入文件不存在: {input_file}")
|
||
return
|
||
|
||
# 生成输出文件名
|
||
output_filename = input_file.stem + args.output_suffix + '.pkl'
|
||
output_file = output_dir / output_filename
|
||
|
||
try:
|
||
# 执行转换
|
||
converted_data = process_phoneme_dataset(input_file, output_file)
|
||
|
||
# 分析结果
|
||
analyze_conversion_results(output_file)
|
||
|
||
print(f"\n转换完成！结果已保存到: {output_file}")
|
||
|
||
except Exception as e:
|
||
print(f"处理过程中出现错误: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
if __name__ == "__main__":
|
||
main() |