File size: 4,805 Bytes
875bba7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
import torch
from torch.utils.data import Dataset
import numpy as np
from tqdm import tqdm
class BeatTrackingDataset(Dataset):
def __init__(
self,
hf_dataset,
target_type="beats",
sample_rate=16000,
hop_length=160,
context_frames=50,
):
"""
Args:
hf_dataset: HuggingFace dataset object
target_type (str): "beats" or "downbeats". Determines which labels are treated as positive.
context_frames (int): Number of frames before and after the center frame.
Total frames = 2 * context_frames + 1.
Default 50 means 101 frames (~1s).
"""
self.sr = sample_rate
self.hop_length = hop_length
self.target_type = target_type
self.context_frames = context_frames
# Context window size in samples
# We need enough samples for the center frame +/- context frames
# PLUS the window size of the largest FFT to compute the edges correctly.
# Largest window in MultiViewSpectrogram is 1488.
self.context_samples = (self.context_frames * 2 + 1) * hop_length + 1488
# Cache audio arrays in memory for fast access
self.audio_cache = []
self.indices = []
self._prepare_indices(hf_dataset)
def _prepare_indices(self, hf_dataset):
"""
Prepares balanced indices and caches audio.
Uses the same "Fuzzier" training examples strategy as the baseline.
"""
print(f"Preparing dataset indices for target: {self.target_type}...")
for i, item in tqdm(
enumerate(hf_dataset), total=len(hf_dataset), desc="Building indices"
):
# Cache audio array (convert to numpy if tensor)
audio = item["audio"]["array"]
if hasattr(audio, "numpy"):
audio = audio.numpy()
self.audio_cache.append(audio)
# Calculate total frames available in audio
audio_len = len(audio)
n_frames = int(audio_len / self.hop_length)
# Select ground truth based on target_type
if self.target_type == "downbeats":
gt_times = item["downbeats"]
else:
gt_times = item["beats"]
# Convert to list if tensor
if hasattr(gt_times, "tolist"):
gt_times = gt_times.tolist()
gt_frames = set([int(t * self.sr / self.hop_length) for t in gt_times])
# --- Positive Examples (with Fuzziness) ---
pos_frames = set()
for bf in gt_frames:
if 0 <= bf < n_frames:
self.indices.append((i, bf, 1.0)) # Center frame
pos_frames.add(bf)
# Neighbors weighted at 0.25
if 0 <= bf - 1 < n_frames:
self.indices.append((i, bf - 1, 0.25))
pos_frames.add(bf - 1)
if 0 <= bf + 1 < n_frames:
self.indices.append((i, bf + 1, 0.25))
pos_frames.add(bf + 1)
# --- Negative Examples ---
# Balance 2:1
num_pos = len(pos_frames)
num_neg = num_pos * 2
count = 0
attempts = 0
while count < num_neg and attempts < num_neg * 5:
f = np.random.randint(0, n_frames)
if f not in pos_frames:
self.indices.append((i, f, 0.0))
count += 1
attempts += 1
print(
f"Dataset ready. {len(self.indices)} samples, {len(self.audio_cache)} tracks cached."
)
def __len__(self):
return len(self.indices)
def __getitem__(self, idx):
track_idx, frame_idx, label = self.indices[idx]
# Fast lookup from cache
audio = self.audio_cache[track_idx]
audio_len = len(audio)
# Calculate sample range for context window
center_sample = frame_idx * self.hop_length
half_context = self.context_samples // 2
# We want the window centered around center_sample
start = center_sample - half_context
end = center_sample + half_context
# Handle padding if needed
pad_left = max(0, -start)
pad_right = max(0, end - audio_len)
valid_start = max(0, start)
valid_end = min(audio_len, end)
# Extract audio chunk
chunk = audio[valid_start:valid_end]
if pad_left > 0 or pad_right > 0:
chunk = np.pad(chunk, (pad_left, pad_right), mode="constant")
waveform = torch.tensor(chunk, dtype=torch.float32)
return waveform, torch.tensor([label], dtype=torch.float32)
|