50 lines
2.1 KiB
Python
50 lines
2.1 KiB
Python
import torch
|
|
from torch.utils.data import Dataset
|
|
import torchaudio
|
|
import os
|
|
|
|
class AudioDataset(Dataset):
|
|
def __init__(self, input_dir, target_duration=None, padding_mode='constant', padding_value=0.0):
|
|
self.input_files = [os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.endswith('.wav')]
|
|
self.target_duration = target_duration # Duration in seconds or None if not set
|
|
self.padding_mode = padding_mode
|
|
self.padding_value = padding_value
|
|
|
|
def __len__(self):
|
|
return len(self.input_files)
|
|
|
|
def __getitem__(self, idx):
|
|
# Load audio samples using torchaudio
|
|
high_quality_wav, sr_original = torchaudio.load(self.input_files[idx], normalize=True)
|
|
|
|
# Resample to 16000 Hz if necessary
|
|
resample_transform = torchaudio.transforms.Resample(sr_original, 16000)
|
|
low_quality_wav = resample_transform(high_quality_wav)
|
|
|
|
# Calculate target length in samples if target_duration is specified
|
|
if self.target_duration is not None:
|
|
target_length = int(self.target_duration * 16000) # Assuming 16000 Hz as target sample rate
|
|
else:
|
|
target_length = high_quality_wav.size(1)
|
|
|
|
# Pad high_quality_wav and low_quality_wav to target_length
|
|
high_quality_wav = self.pad_tensor(high_quality_wav, target_length)
|
|
low_quality_wav = self.pad_tensor(low_quality_wav, target_length)
|
|
|
|
return high_quality_wav, low_quality_wav
|
|
|
|
def pad_tensor(self, tensor, target_length):
|
|
"""Pad tensor to target length along the time dimension (dim=1)."""
|
|
current_length = tensor.size(1)
|
|
|
|
if current_length < target_length:
|
|
# Calculate padding amount for each side
|
|
padding_amount = target_length - current_length
|
|
padding = (0, padding_amount) # (left_pad, right_pad) for 1D padding
|
|
tensor = torch.nn.functional.pad(tensor, padding, mode=self.padding_mode, value=self.padding_value)
|
|
else:
|
|
# If tensor is longer than target, truncate it
|
|
tensor = tensor[:, :target_length]
|
|
|
|
return tensor
|