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 |