File size: 4,643 Bytes
707cbac |
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 |
import torch
from torch.utils.data import Dataset
import numpy as np
from tqdm import tqdm
from .utils import extract_context
class BeatTrackingDataset(Dataset):
def __init__(
self, hf_dataset, target_type="beats", sample_rate=16000, hop_length=160
):
"""
Args:
hf_dataset: HuggingFace dataset object
target_type (str): "beats" or "downbeats". Determines which labels are treated as positive.
"""
self.sr = sample_rate
self.hop_length = hop_length
self.target_type = target_type
# Context window size in samples (7 frames = 70ms at 100fps)
self.context_frames = 7
self.context_samples = (self.context_frames * 2 + 1) * hop_length + max(
[368, 736, 1488]
) # extra for FFT window
# 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.
Paper Section 4.5: Uses "Fuzzier" training examples (neighbors weighted less).
"""
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":
# Only downbeats are positives
gt_times = item["downbeats"]
else:
# All beats are positives (downbeats are also beats)
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) ---
# "define a single frame before and after each annotated onset to be additional positive examples"
pos_frames = set()
for bf in gt_frames:
if 0 <= bf < n_frames:
self.indices.append((i, bf, 1.0)) # Center frame (Sharp onset)
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 ---
# Paper uses "all others as negative", but we balance 2:1 for stable SGD.
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
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)
start = max(0, start)
end = min(audio_len, end)
# Extract audio chunk
chunk = audio[start: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)
|