add utils

utils.py

@@ -0,0 +1,226 @@
+# Ke Chen
+# knutchen@ucsd.edu
+# HTS-AT: A HIERARCHICAL TOKEN-SEMANTIC AUDIO TRANSFORMER FOR SOUND CLASSIFICATION AND DETECTION
+# Some Useful Common Methods
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch import Tensor
+from typing import Optional
+import logging
+import os
+import sys
+import h5py
+import csv
+import time
+import json
+import museval
+import librosa
+from datetime import datetime
+from tqdm import tqdm
+from scipy import stats
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# import from https://github.com/Alibaba-MIIL/ASL/blob/main/src/loss_functions/losses.py
+class AsymmetricLoss(nn.Module):
+    def __init__(self, gamma_neg=4, gamma_pos=1, clip=0.05, eps=1e-8, disable_torch_grad_focal_loss=True):
+        super(AsymmetricLoss, self).__init__()
+
+        self.gamma_neg = gamma_neg
+        self.gamma_pos = gamma_pos
+        self.clip = clip
+        self.disable_torch_grad_focal_loss = disable_torch_grad_focal_loss
+        self.eps = eps
+
+    def forward(self, x, y):
+        """"
+        Parameters
+        ----------
+        x: input logits
+        y: targets (multi-label binarized vector)
+        """
+
+        # Calculating Probabilities
+        # x_sigmoid = torch.sigmoid(x)
+        x_sigmoid = x # without sigmoid since it has been computed
+        xs_pos = x_sigmoid
+        xs_neg = 1 - x_sigmoid
+
+        # Asymmetric Clipping
+        if self.clip is not None and self.clip > 0:
+            xs_neg = (xs_neg + self.clip).clamp(max=1)
+
+        # Basic CE calculation
+        los_pos = y * torch.log(xs_pos.clamp(min=self.eps))
+        los_neg = (1 - y) * torch.log(xs_neg.clamp(min=self.eps))
+        loss = los_pos + los_neg
+
+        # Asymmetric Focusing
+        if self.gamma_neg > 0 or self.gamma_pos > 0:
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(False)
+            pt0 = xs_pos * y
+            pt1 = xs_neg * (1 - y)  # pt = p if t > 0 else 1-p
+            pt = pt0 + pt1
+            one_sided_gamma = self.gamma_pos * y + self.gamma_neg * (1 - y)
+            one_sided_w = torch.pow(1 - pt, one_sided_gamma)
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(True)
+            loss *= one_sided_w
+
+        return -loss.mean()
+
+
+def get_mix_lambda(mixup_alpha, batch_size):
+    mixup_lambdas = [np.random.beta(mixup_alpha, mixup_alpha, 1)[0] for _ in range(batch_size)]
+    return np.array(mixup_lambdas).astype(np.float32)
+
+def create_folder(fd):
+    if not os.path.exists(fd):
+        os.makedirs(fd)
+
+def dump_config(config, filename, include_time = False):
+    save_time = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
+    config_json = {}
+    for key in dir(config):
+        if not key.startswith("_"):
+            config_json[key] = eval("config." + key)
+    if include_time:
+        filename = filename + "_" + save_time
+    with open(filename + ".json", "w") as f:      
+        json.dump(config_json, f ,indent=4)
+
+def int16_to_float32(x):
+    return (x / 32767.).astype(np.float32)
+
+def float32_to_int16(x):
+    x = np.clip(x, a_min = -1., a_max = 1.)
+    return (x * 32767.).astype(np.int16)
+
+
+# index for each class
+def process_idc(index_path, classes_num, filename):
+    # load data
+    logging.info("Load Data...............")
+    idc = [[] for _ in range(classes_num)]
+    with h5py.File(index_path, "r") as f:
+        for i in tqdm(range(len(f["target"]))):
+            t_class = np.where(f["target"][i])[0]
+            for t in t_class:
+                idc[t].append(i)
+    print(idc)
+    np.save(filename, idc)
+    logging.info("Load Data Succeed...............")
+
+def clip_bce(pred, target):
+    """Binary crossentropy loss.
+    """
+    return F.cross_entropy(pred, target)
+    # return F.binary_cross_entropy(pred, target)
+
+
+def clip_ce(pred, target):
+    return F.cross_entropy(pred, target)
+
+def d_prime(auc):
+    d_prime = stats.norm().ppf(auc) * np.sqrt(2.0)
+    return d_prime
+
+
+def get_loss_func(loss_type):
+    if loss_type == 'clip_bce':
+        return clip_bce
+    if loss_type == 'clip_ce':
+        return clip_ce
+    if loss_type == 'asl_loss':
+        loss_func = AsymmetricLoss(gamma_neg=4, gamma_pos=0,clip=0.05)
+        return loss_func
+
+def do_mixup_label(x):
+    out = torch.logical_or(x, torch.flip(x, dims = [0])).float()
+    return out
+
+def do_mixup(x, mixup_lambda):
+    """
+    Args:
+      x: (batch_size , ...)
+      mixup_lambda: (batch_size,)
+
+    Returns:
+      out: (batch_size, ...)
+    """
+    out = (x.transpose(0,-1) * mixup_lambda + torch.flip(x, dims = [0]).transpose(0,-1) * (1 - mixup_lambda)).transpose(0,-1)
+    return out
+    
+def interpolate(x, ratio):
+    """Interpolate data in time domain. This is used to compensate the 
+    resolution reduction in downsampling of a CNN.
+    
+    Args:
+      x: (batch_size, time_steps, classes_num)
+      ratio: int, ratio to interpolate
+
+    Returns:
+      upsampled: (batch_size, time_steps * ratio, classes_num)
+    """
+    (batch_size, time_steps, classes_num) = x.shape
+    upsampled = x[:, :, None, :].repeat(1, 1, ratio, 1)
+    upsampled = upsampled.reshape(batch_size, time_steps * ratio, classes_num)
+    return upsampled
+
+
+def pad_framewise_output(framewise_output, frames_num):
+    """Pad framewise_output to the same length as input frames. The pad value 
+    is the same as the value of the last frame.
+
+    Args:
+      framewise_output: (batch_size, frames_num, classes_num)
+      frames_num: int, number of frames to pad
+
+    Outputs:
+      output: (batch_size, frames_num, classes_num)
+    """
+    pad = framewise_output[:, -1 :, :].repeat(1, frames_num - framewise_output.shape[1], 1)
+    """tensor for padding"""
+
+    output = torch.cat((framewise_output, pad), dim=1)
+    """(batch_size, frames_num, classes_num)"""
+
+    return output
+
+# set the audio into the format that can be fed into the model
+# resample -> convert to mono -> output the audio  
+# track [n_sample, n_channel]
+def prepprocess_audio(track, ofs, rfs, mono_type = "mix"):
+    if track.shape[-1] > 1:
+        # stereo
+        if mono_type == "mix":
+            track = np.transpose(track, (1,0))
+            track = librosa.to_mono(track)
+        elif mono_type == "left":
+            track = track[:, 0]
+        elif mono_type == "right":
+            track = track[:, 1]
+    else:
+        track = track[:, 0]
+    # track [n_sample]
+    if ofs != rfs:
+        track = librosa.resample(track, ofs, rfs)
+    return track
+
+def init_hier_head(class_map, num_class):
+    class_map = np.load(class_map, allow_pickle = True)
+    
+    head_weight = torch.zeros(num_class,num_class).float()
+    head_bias = torch.zeros(num_class).float()
+
+    for i in range(len(class_map)):
+        for d in class_map[i][1]:
+            head_weight[d][i] = 1.0
+        for d in class_map[i][2]:
+            head_weight[d][i] = 1.0 / len(class_map[i][2])
+        head_weight[i][i] = 1.0
+    return head_weight, head_bias