Spaces:
Sleeping
Sleeping
File size: 2,282 Bytes
5e740f6 |
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 |
import os
import torch
from .pytorch_utils import move_data_to_device
from .models import Cnn14_DecisionLevelMax
from .config import labels, classes_num
def create_folder(fd):
if not os.path.exists(fd):
os.makedirs(fd)
def get_filename(path):
path = os.path.realpath(path)
na_ext = path.split('/')[-1]
na = os.path.splitext(na_ext)[0]
return na
class SoundEventDetection(object):
def __init__(self, model=None, checkpoint_path=None, device='cuda'):
"""Sound event detection inference wrapper.
"""
if not checkpoint_path:
checkpoint_path='panns_data/Cnn14_DecisionLevelMax.pth' # moved to current directory
print('Checkpoint path: {}'.format(checkpoint_path))
if not os.path.exists(checkpoint_path) or os.path.getsize(checkpoint_path) < 3e8:
create_folder(os.path.dirname(checkpoint_path))
os.system('wget -O "{}" https://zenodo.org/record/3987831/files/Cnn14_DecisionLevelMax_mAP%3D0.385.pth?download=1'.format(checkpoint_path))
if device == 'cuda' and torch.cuda.is_available():
self.device = 'cuda'
else:
self.device = 'cpu'
self.labels = labels
self.classes_num = classes_num
# Model
if model is None:
self.model = Cnn14_DecisionLevelMax(sample_rate=32000, window_size=1024,
hop_size=320, mel_bins=64, fmin=50, fmax=14000,
classes_num=self.classes_num)
else:
self.model = model
checkpoint = torch.load(checkpoint_path, map_location=self.device)
self.model.load_state_dict(checkpoint['model'])
# Parallel
if 'cuda' in str(self.device):
self.model.to(self.device)
print('GPU number: {}'.format(torch.cuda.device_count()))
self.model = torch.nn.DataParallel(self.model)
else:
print('Using CPU.')
def inference(self, audio):
audio = move_data_to_device(audio, self.device)
with torch.no_grad():
self.model.eval()
output_dict = self.model(audio, None)
framewise_output = output_dict['framewise_output'].data.cpu().numpy()
return framewise_output
|