본문 바로가기
코딩

TTS AI 예제 3

by 나는야석사 2024. 7. 11.
import os
import json
import numpy as np
import librosa
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, Trainer, TrainingArguments
from datasets import Dataset
from sklearn.model_selection import train_test_split
import evaluate
import matplotlib.pyplot as plt

# 음성 파일 및 라벨 데이터 폴더 경로 설정
audio_dir = '/content/path_to_audio_folder'  # Colab에서는 content 폴더 내에 경로 설정
label_dir = '/content/path_to_label_folder'  # Colab에서는 content 폴더 내에 경로 설정

# 파일 확장자가 .wav인 파일 목록 가져오기
audio_files = [f for f in os.listdir(audio_dir) if f.endswith('.wav')]

# 파일 확장자가 .json인 라벨 파일 목록 가져오기
label_files = [f for f in os.listdir(label_dir) if f.endswith('.json')]

# 음성 파일을 로드하고 전처리 함수 정의
def load_audio_file(file_path, sr=16000):
    y, sr = librosa.load(file_path, sr=sr)
    return y, sr

# 모든 라벨 파일 로드 및 병합
all_labels = {}
for label_file in label_files:
    with open(os.path.join(label_dir, label_file), 'r', encoding='utf-8') as f:
        labels = json.load(f)
        file_name = labels['파일정보']['FileName']
        all_labels[file_name] = labels

# 데이터 준비
X = []
y = []

for audio_file in audio_files:
    file_path = os.path.join(audio_dir, audio_file)
    if audio_file in all_labels:
        y_audio, sr = load_audio_file(file_path)
        label_data = all_labels[audio_file]
        if "OrgLabelText" in label_data["전사정보"]:
            transcript = label_data["전사정보"]["OrgLabelText"]
            X.append({"array": y_audio, "sampling_rate": sr})
            y.append(transcript)

# X, y 출력 확인
print("X: ", len(X))  # 데이터 개수 확인
print("y: ", len(y))  # 데이터 개수 확인

# train_test_split을 사용하여 데이터 분할
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 데이터셋 준비
train_data = Dataset.from_list([{"audio": X, "transcript": y} for X, y in zip(X_train, y_train)])
test_data = Dataset.from_list([{"audio": X, "transcript": y} for X, y in zip(X_test, y_test)])

# Wav2Vec 2.0 모델 및 프로세서 로드
processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-960h")
model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-960h")

# 데이터셋 전처리
def prepare_dataset(batch):
    audio = batch["audio"]
    inputs = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt", padding="max_length", truncation=True, max_length=32000)
    labels = processor(text=[batch["transcript"]], return_tensors="pt", padding="max_length", truncation=True, max_length=128).input_ids
    batch["input_values"] = inputs.input_values[0].numpy()  # Correctly reshape input_values
    batch["labels"] = labels[0].numpy()
    return batch

# 데이터셋 처리
train_dataset = train_data.map(prepare_dataset, remove_columns=["audio", "transcript"])
test_dataset = test_data.map(prepare_dataset, remove_columns=["audio", "transcript"])

# 데이터셋 크기 확인
print(f"Train dataset size: {len(train_dataset)}")
print(f"Test dataset size: {len(test_dataset)}")

# 데이터셋의 일부를 출력하여 확인
for i in range(5):
    sample = train_dataset[i]
    print(f"Sample {i+1}:")
    print(f"  Input values: {sample['input_values'][:10]}")  # 입력 값의 처음 10개 값만 출력
    print(f"  Labels: {sample['labels']}")
    print()

def plot_waveform_and_transcript(sample, index):
    plt.figure(figsize=(14, 5))
    plt.plot(sample['input_values'])
    plt.title(f"Sample {index+1} - Transcript: {processor.decode(sample['labels'])}")
    plt.xlabel("Time")
    plt.ylabel("Amplitude")
    plt.show()

# 데이터셋의 일부 샘플을 시각화하여 확인
for i in range(5):
    sample = train_dataset[i]
    plot_waveform_and_transcript(sample, i)

# 모델 학습 설정
training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    save_steps=10,
    save_total_limit=2,
    logging_dir='./logs',  # 로그 디렉토리 추가
)

# 모델 평가 함수 정의
wer_metric = evaluate.load("wer")

def compute_metrics(pred):
    pred_logits = pred.predictions
    pred_ids = np.argmax(pred_logits, axis=-1)
    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
    labels_ids = pred.label_ids
    labels_str = processor.batch_decode(labels_ids, skip_special_tokens=True)

    # 정확도 계산
    correct_predictions = sum([1 for pred, label in zip(pred_str, labels_str) if pred == label])
    accuracy = correct_predictions / len(pred_str)

    # WER 계산
    wer = wer_metric.compute(predictions=pred_str, references=labels_str)

    return {"accuracy": accuracy, "wer": wer}

# 트레이너 설정
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    tokenizer=processor.feature_extractor,
    compute_metrics=compute_metrics,  # compute_metrics 함수 추가
)

# 모델 학습
trainer.train()

# 모델 평가
results = trainer.evaluate()
print(results)  # 반환된 결과 구조를 확인
print(f"Accuracy: {results['eval_accuracy'] * 100:.2f}%")
print(f"Word Error Rate (WER): {results['eval_wer'] * 100:.2f}%")

 

 

<unknown> 아마 이게 더 정확?

import os
import json
import numpy as np
import librosa
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, Trainer, TrainingArguments
from datasets import Dataset
from sklearn.model_selection import train_test_split
import evaluate
import matplotlib.pyplot as plt

# 음성 파일 및 라벨 데이터 폴더 경로 설정
audio_dir = '/content/path_to_audio_folder'  # Colab에서는 content 폴더 내에 경로 설정
label_dir = '/content/path_to_label_folder'  # Colab에서는 content 폴더 내에 경로 설정

# 파일 확장자가 .wav인 파일 목록 가져오기
audio_files = [f for f in os.listdir(audio_dir) if f.endswith('.wav')]

# 파일 확장자가 .json인 라벨 파일 목록 가져오기
label_files = [f for f in os.listdir(label_dir) if f.endswith('.json')]

# 음성 파일을 로드하고 전처리 함수 정의
def load_audio_file(file_path, sr=16000):
    y, sr = librosa.load(file_path, sr=sr)
    return y, sr

# 모든 라벨 파일 로드 및 병합
all_labels = {}
for label_file in label_files:
    with open(os.path.join(label_dir, label_file), 'r', encoding='utf-8') as f:
        labels = json.load(f)
        file_name = labels['파일정보']['FileName']
        all_labels[file_name] = labels

# 데이터 준비
X = []
y = []

for audio_file in audio_files:
    file_path = os.path.join(audio_dir, audio_file)
    if audio_file in all_labels:
        y_audio, sr = load_audio_file(file_path)
        label_data = all_labels[audio_file]
        if "OrgLabelText" in label_data["전사정보"]:
            transcript = label_data["전사정보"]["OrgLabelText"]
            X.append({"array": y_audio, "sampling_rate": sr})
            y.append(transcript)

# X, y 출력 확인
print("X: ", len(X))  # 데이터 개수 확인
print("y: ", len(y))  # 데이터 개수 확인

# train_test_split을 사용하여 데이터 분할
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 데이터셋 준비
train_data = Dataset.from_list([{"audio": X, "transcript": y} for X, y in zip(X_train, y_train)])
test_data = Dataset.from_list([{"audio": X, "transcript": y} for X, y in zip(X_test, y_test)])

# Wav2Vec 2.0 모델 및 프로세서 로드
processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-960h")
model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-960h")

# 데이터셋 전처리
def prepare_dataset(batch):
    audio = batch["audio"]
    inputs = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt", padding="max_length", truncation=True, max_length=32000)
    labels = processor(text=[batch["transcript"]], return_tensors="pt", padding="max_length", truncation=True, max_length=128).input_ids
    batch["input_values"] = inputs.input_values[0].numpy()  # Correctly reshape input_values
    batch["labels"] = labels[0].numpy()
    return batch

# 데이터셋 처리
train_dataset = train_data.map(prepare_dataset, remove_columns=["audio", "transcript"])
test_dataset = test_data.map(prepare_dataset, remove_columns=["audio", "transcript"])

# 데이터셋 크기 확인
print(f"Train dataset size: {len(train_dataset)}")
print(f"Test dataset size: {len(test_dataset)}")

# 데이터셋의 일부를 출력하여 확인
for i in range(5):
    sample = train_dataset[i]
    print(f"Sample {i+1}:")
    print(f"  Input values: {sample['input_values'][:10]}")  # 입력 값의 처음 10개 값만 출력
    print(f"  Labels: {sample['labels']}")
    print()

# 사전 업데이트
def update_vocab(processor, transcripts):
    tokenizer = processor.tokenizer
    new_tokens = set()
    for transcript in transcripts:
        new_tokens.update(tokenizer.tokenize(transcript))
    tokenizer.add_tokens(list(new_tokens))
    processor.feature_extractor._processor = processor
    return processor

# 업데이트된 사전으로 프로세서 초기화
processor = update_vocab(processor, y)

def plot_waveform_and_transcript(sample, index):
    plt.figure(figsize=(14, 5))
    plt.plot(sample['input_values'])
    decoded_text = processor.tokenizer.decode(sample['labels'], skip_special_tokens=True)
    plt.title(f"Sample {index+1} - Transcript: {decoded_text}")
    plt.xlabel("Time")
    plt.ylabel("Amplitude")
    plt.show()

# 데이터셋의 일부 샘플을 시각화하여 확인
for i in range(5):
    sample = train_dataset[i]
    plot_waveform_and_transcript(sample, i)

# 모델 학습 설정
training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    save_steps=10,
    save_total_limit=2,
    logging_dir='./logs',  # 로그 디렉토리 추가
)

# 모델 평가 함수 정의
wer_metric = evaluate.load("wer")

def compute_metrics(pred):
    pred_logits = pred.predictions
    pred_ids = np.argmax(pred_logits, axis=-1)
    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
    labels_ids = pred.label_ids
    labels_str = processor.batch_decode(labels_ids, skip_special_tokens=True)

    # 정확도 계산
    correct_predictions = sum([1 for pred, label in zip(pred_str, labels_str) if pred == label])
    accuracy = correct_predictions / len(pred_str)

    # WER 계산
    wer = wer_metric.compute(predictions=pred_str, references=labels_str)

    return {"accuracy": accuracy, "wer": wer}

# 트레이너 설정
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    tokenizer=processor.feature_extractor,
    compute_metrics=compute_metrics,  # compute_metrics 함수 추가
)

# 모델 학습
trainer.train()

# 모델 평가
results = trainer.evaluate()
print(results)  # 반환된 결과 구조를 확인
print(f"Accuracy: {results['eval_accuracy'] * 100:.2f}%")
print(f"Word Error Rate (WER): {results['eval_wer'] * 100:.2f}%")

'코딩' 카테고리의 다른 글

TTS AI 예제 여러 개  (0) 2024.07.11
여러 개의 wav 파일을 하나로 합치는 방법  (0) 2024.07.11
TTS AI 예제 2  (0) 2024.07.11
TTS AI 예제 1  (0) 2024.07.11
Git 설치 방법  (0) 2024.04.10