Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import numpy as np | |
| np.seterr(divide='ignore', invalid='ignore') | |
| def x_coords_to_time(x_pos, sampling_rate, fft_win_length, fft_overlap): | |
| nfft = int(fft_win_length*sampling_rate) | |
| noverlap = int(fft_overlap*nfft) | |
| return ((x_pos*(nfft - noverlap)) + noverlap) / sampling_rate | |
| #return (1.0 - fft_overlap) * fft_win_length * (x_pos + 0.5) # 0.5 is for center of temporal window | |
| def overall_class_pred(det_prob, class_prob): | |
| weighted_pred = (class_prob*det_prob).sum(1) | |
| return weighted_pred / weighted_pred.sum() | |
| def run_nms(outputs, params, sampling_rate): | |
| pred_det = outputs['pred_det'] # probability of box | |
| pred_size = outputs['pred_size'] # box size | |
| pred_det_nms = non_max_suppression(pred_det, params['nms_kernel_size']) | |
| freq_rescale = (params['max_freq'] - params['min_freq']) /pred_det.shape[-2] | |
| # NOTE there will be small differences depending on which sampling rate is chosen | |
| # as we are choosing the same sampling rate for the entire batch | |
| duration = x_coords_to_time(pred_det.shape[-1], sampling_rate[0].item(), | |
| params['fft_win_length'], params['fft_overlap']) | |
| top_k = int(duration * params['nms_top_k_per_sec']) | |
| scores, y_pos, x_pos = get_topk_scores(pred_det_nms, top_k) | |
| # loop over batch to save outputs | |
| preds = [] | |
| feats = [] | |
| for ii in range(pred_det_nms.shape[0]): | |
| # get valid indices | |
| inds_ord = torch.argsort(x_pos[ii, :]) | |
| valid_inds = scores[ii, inds_ord] > params['detection_threshold'] | |
| valid_inds = inds_ord[valid_inds] | |
| # create result dictionary | |
| pred = {} | |
| pred['det_probs'] = scores[ii, valid_inds] | |
| pred['x_pos'] = x_pos[ii, valid_inds] | |
| pred['y_pos'] = y_pos[ii, valid_inds] | |
| pred['bb_width'] = pred_size[ii, 0, pred['y_pos'], pred['x_pos']] | |
| pred['bb_height'] = pred_size[ii, 1, pred['y_pos'], pred['x_pos']] | |
| pred['start_times'] = x_coords_to_time(pred['x_pos'].float() / params['resize_factor'], | |
| sampling_rate[ii].item(), params['fft_win_length'], params['fft_overlap']) | |
| pred['end_times'] = x_coords_to_time((pred['x_pos'].float()+pred['bb_width']) / params['resize_factor'], | |
| sampling_rate[ii].item(), params['fft_win_length'], params['fft_overlap']) | |
| pred['low_freqs'] = (pred_size[ii].shape[1] - pred['y_pos'].float())*freq_rescale + params['min_freq'] | |
| pred['high_freqs'] = pred['low_freqs'] + pred['bb_height']*freq_rescale | |
| # extract the per class votes | |
| if 'pred_class' in outputs: | |
| pred['class_probs'] = outputs['pred_class'][ii, :, y_pos[ii, valid_inds], x_pos[ii, valid_inds]] | |
| # extract the model features | |
| if 'features' in outputs: | |
| feat = outputs['features'][ii, :, y_pos[ii, valid_inds], x_pos[ii, valid_inds]].transpose(0, 1) | |
| feat = feat.cpu().numpy().astype(np.float32) | |
| feats.append(feat) | |
| # convert to numpy | |
| for kk in pred.keys(): | |
| pred[kk] = pred[kk].cpu().numpy().astype(np.float32) | |
| preds.append(pred) | |
| return preds, feats | |
| def non_max_suppression(heat, kernel_size): | |
| # kernel can be an int or list/tuple | |
| if type(kernel_size) is int: | |
| kernel_size_h = kernel_size | |
| kernel_size_w = kernel_size | |
| pad_h = (kernel_size_h - 1) // 2 | |
| pad_w = (kernel_size_w - 1) // 2 | |
| hmax = nn.functional.max_pool2d(heat, (kernel_size_h, kernel_size_w), stride=1, padding=(pad_h, pad_w)) | |
| keep = (hmax == heat).float() | |
| return heat * keep | |
| def get_topk_scores(scores, K): | |
| # expects input of size: batch x 1 x height x width | |
| batch, _, height, width = scores.size() | |
| topk_scores, topk_inds = torch.topk(scores.view(batch, -1), K) | |
| topk_inds = topk_inds % (height * width) | |
| topk_ys = torch.div(topk_inds, width, rounding_mode='floor').long() | |
| topk_xs = (topk_inds % width).long() | |
| return topk_scores, topk_ys, topk_xs | |