| import os, sys, torch, joblib, importlib |
| import pandas as pd |
| import torchaudio.transforms as T |
|
|
| class BioacousticEngine: |
| def __init__(self, repo_dir="tiny-bird-diffusion"): |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self.repo_dir = repo_dir |
| sys.path.append(os.path.abspath(repo_dir)) |
| from cvt import cvt13 |
| mel_module = importlib.import_module("mel_spectrogram") |
| self.preprocessor = mel_module.MelSpectrogramProcessor(device=self.device) |
| self.model = cvt13() |
| self.model.load_state_dict(torch.load(f"{repo_dir}/protoclr.pth", map_location="cpu")) |
| self.model = self.model.to(self.device).eval() |
| brain_data = joblib.load(f"{repo_dir}/trained_cluster_brain.joblib") |
| self.reducer = brain_data['umap'] |
| self.df = pd.read_csv(f"{repo_dir}/acoustic_atlas_metadata.csv") |
|
|
| def process_waveform(self, waveform, sample_rate): |
| if sample_rate != 16000: waveform = T.Resample(orig_freq=sample_rate, new_freq=16000)(waveform) |
| if waveform.shape[0] > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) |
| total_samples = waveform.shape[-1] |
| target_samples = 3 * 16000 |
| if total_samples > target_samples: |
| step, max_energy, best_start = 4000, -1, 0 |
| for start in range(0, total_samples - target_samples + 1, step): |
| energy = waveform[:, start:start + target_samples].abs().mean().item() |
| if energy > max_energy: max_energy, best_start = energy, start |
| waveform = waveform[:, best_start:best_start + target_samples] |
| if waveform.abs().max() > 0.02: waveform = waveform / waveform.abs().max() |
| return waveform.to(self.device) |
|
|
