from typing import Any, Dict, Tuple import torch import torch.nn.functional as F import numpy as np from collections import defaultdict, Counter import heapq SAMPLE_RATE = 16000 WINDOW_SIZE = 512 def chunked(iterable, size): """Helper to yield successive chunks of given size from iterable.""" for i in range(0, len(iterable), size): yield iterable[i : i + size] class GPUShazam: def __init__(self, sample_rate=SAMPLE_RATE, n_fft=WINDOW_SIZE, hop_length=None, fan_value=5, min_match_ratio: float = 0.1, device='cuda', kernal_size=5, min_prune_size=500, max_prune_ratio=0.8): """ GPU-based Shazam with optional query‐match threshold. Args: … min_query_matches: only accept a match if vote count ≥ this threshold """ self.sample_rate = sample_rate self.n_fft = n_fft self.hop_length = hop_length or n_fft // 4 self.fan_value = fan_value self.min_match_ratio = min_match_ratio self.window = torch.hann_window(self.n_fft, device=device) self.index = {} self.min_hash_time_delta = 0 self.max_hash_time_delta = 200 self.min_hash_freq_delta = -30 self.max_hash_freq_delta = 30 self.kernal_size = kernal_size self.device = device self.min_prune_size = min_prune_size self.max_prune_ratio = max_prune_ratio self._set = set() def fingerprint(self, waveform, delta_compress: bool = False): """ Compute fingerprints for a single audio waveform. Args: waveform (1D array-like): audio samples Returns: List of (hash, time_offset) tuples. """ # Move signal to GPU sig = torch.tensor(waveform, dtype=torch.float32, device=self.device) # Compute complex STFT on GPU spec = torch.stft(sig, n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.n_fft, window=self.window, center=False, return_complex=True) # Magnitude spectrogram mag = spec.abs() # Local max pooling for peak detection padding = self.kernal_size // 2 max_pooled = F.max_pool2d( mag.unsqueeze(0).unsqueeze(0), kernel_size=(self.kernal_size, self.kernal_size), stride=1, padding=(padding, padding) ) #amp_min = torch.quantile(mag.flatten(), 0.99).item() amp_min = 10 peaks = (mag.unsqueeze(0).unsqueeze(0) == max_pooled) & (mag.unsqueeze(0).unsqueeze(0) > amp_min) peaks = peaks.squeeze().cpu().numpy() freqs, times = np.where(peaks) # after you compute freqs, times: peak_list = list(zip(times, freqs)) if len(peak_list) == 0: return [] peak_list.sort(key=lambda x: x[0]) # sort by time times_sorted, freqs_sorted = zip(*peak_list) # now pair in temporal order hashes = [] for i in range(len(times_sorted)): for j in range(1, self.fan_value): if i + j < len(times_sorted): t1, f1 = times_sorted[i], freqs_sorted[i] t2, f2 = times_sorted[i+j], freqs_sorted[i+j] dt = t2 - t1 df = f2 - f1 if self.min_hash_time_delta <= dt <= self.max_hash_time_delta and self.min_hash_freq_delta <= df <= self.max_hash_freq_delta: h = (f1 << 24) | (f2 << 16) | dt hashes.append((h, t1)) if not delta_compress: return hashes # 3) aggregate into a list of (h, [t1, t2, …]) grouping = defaultdict(list) for h, t in hashes: grouping[h].append(t) return [(h, grouping[h]) for h in grouping] def add_song(self, song_id, waveform): """ Add a song to the index. Args: song_id (Any): unique identifier for the song waveform (1D array-like): audio samples """ fpt= self.fingerprint(waveform, delta_compress=True) if len(fpt)>0: self._set.add(song_id) for h, t in fpt: self.index.setdefault(h, []).append((song_id, t)) def top_candidates(self, batch, top_n): counts : Dict[Tuple[int, int], int] = {} for h, t_query in batch: for song_id, t_ref_list in self.index.get(h, []): # t_ref_list 可能是列表，也可能是单个值 if isinstance(t_ref_list, list): for t_ref in t_ref_list: delta = t_ref - t_query counts.setdefault((song_id, delta), 0) counts[(song_id, delta)] += 1 else: t_ref = t_ref_list delta = t_ref - t_query counts.setdefault((song_id, delta), 0) counts[(song_id, delta)] += 1 counter = Counter(counts) return counter.most_common(top_n) def _prune(self): if len(self._set) < self.min_prune_size: return maximum_cnt = len(self._set) * self.max_prune_ratio for hash in self.index: if len(self.index[hash]) >= maximum_cnt: self.index[hash] = [] def most_common(self, counter_dict, top_n): counter = Counter(counter_dict) return counter.most_common(top_n) def query(self, waveform, top_n=1): """ Returns (best_song_id, best_time_delta, match_ratio). If no match or ratio < min_match_ratio, song_id/time_delta will be None. """ """ 返回匹配度最高的 top_n 个结果列表，每个元素为 (song_id, time_delta, match_ratio)。 Args: waveform: 1D numpy array，待查询的音频片段 top_n : 整数，返回前 top_n 项（如果总候选数不足，也可能少于 top_n）。 返回: List[ (song_id or None, time_delta or None, match_ratio) ] - 如果 match_ratio < self.min_match_ratio，则该项的 song_id/time_delta 均为 None 但仍返回 match_ratio。 """ # 1) 对片段做指纹提取 hashes = list(self.fingerprint(waveform, delta_compress=False)) total_hashes = len(hashes) if total_hashes == 0: # 如果根本没有哈希，就直接返回 top_n 个 (None, None, 0.0) return [(None, None, 0.0)] * top_n # 2) 统计每个 (song_id, Δt) 的匹配数 counts: Dict[(Any, int), int] = {} batch_size = 1000 candidates, total = {}, 0 for batch in chunked(hashes, batch_size): print(f"Processing batches in query, {total + len(batch)}") top_candidates = self.top_candidates(batch, top_n) # 4) 最终拼出结果列表 results, should_stop = [], True for tp, vote_count in top_candidates: original_count = candidates.get(tp, 0) original_ratio = original_count / total if total else 0 new_match_ratio = (original_count + vote_count) / (total + len(batch)) if new_match_ratio <=self.min_match_ratio and abs(new_match_ratio - original_ratio) >= 0.02: should_stop = False candidates[tp] = original_count + vote_count if should_stop: return self.most_common(candidates, top_n), (total + len(batch)) total += len(batch) return self.most_common(candidates, top_n), total