# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import librosa
import os
import numpy as np
import matplotlib.pyplot as plt
from scipy.fftpack import dct
from scipy.signal import spectrogram
import operator
import pickle
import time
import csv
from random import shuffle
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from sklearn import metrics
import logging

# ## for GPU only

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

print(device)


# ### CHOICE 1 : Network without dropout or batch norm
# - comment line `net.to(device)` to run on CPU

# +
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1,128,3)
        self.conv2 = nn.Conv2d(128,384,3)
        self.conv3 = nn.Conv2d(384,768,3)
        self.conv4 = nn.Conv2d(768,2048,3)
        self.conv5 = nn.Conv2d(2048,50,1) # network with 50 outputs
#         self.conv5 = nn.Conv2d(2048,10,1) # network with 10 outputs
        self.pool1 = nn.MaxPool2d((2,4),(2,4))
        self.pool2 = nn.MaxPool2d((2,4),(2,4))
        self.pool3 = nn.MaxPool2d((2,6),(2,6))
        self.pool4 = nn.MaxPool2d((1,7),(1,7))
#         self.sig = nn.Sigmoid()
        
    def forward(self,x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = self.pool3(F.relu(self.conv3(x)))
        x = self.pool4(F.relu(self.conv4(x)))
        x = self.conv5(x)
        x = x.view(-1, 50) # network with 50 outputs
#         x = x.view(-1,10) # network with 10 outputs
#         x = self.sig(x)
        return x
    
net = Net()
net.to(device)

# -

# ### CHOICE 2 : Network with drop out and batch norm
# - comment the line `net.to(device)` to run on CPU

# +
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1,128,3)
        self.norm1 = nn.BatchNorm2d(128)
        self.conv2 = nn.Conv2d(128,384,3)
        self.norm2 = nn.BatchNorm2d(384)
        self.conv3 = nn.Conv2d(384,768,3)
        self.norm3 = nn.BatchNorm2d(768)
        self.conv4 = nn.Conv2d(768,2048,3)
        self.conv5 = nn.Conv2d(2048,50,1) # network with 50 outputs
#         self.conv5 = nn.Conv2d(2048,10,1) # network with 10 outputs
        self.pool1 = nn.MaxPool2d((2,4),(2,4))
        self.pool2 = nn.MaxPool2d((2,4),(2,4))
        self.pool3 = nn.MaxPool2d((2,6),(2,6))
        self.pool4 = nn.MaxPool2d((1,7),(1,7))
        self.drop = nn.Dropout2d(.5)
#         self.sig = nn.Sigmoid()
        
    def forward(self,x):
        x = self.pool1(F.relu(self.norm1(self.conv1(x))))
        x = self.pool2(self.drop(F.relu(self.norm2(self.conv2(x)))))
        x = self.pool3(self.drop(F.relu(self.norm3(self.conv3(x)))))
        x = self.pool4(F.relu(self.conv4(x)))
        x = self.conv5(x)
        x = x.view(-1, 50) # network with 50 outputs
#         x = x.view(-1,10) # network with 10 outputs
#         x = self.sig(x)
        return x
    
net = Net()
net.to(device)


# -

# ### CHOICE 3: new feature space

# +
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv0 = nn.Conv2d(1,40,(40,2))
        self.conv1 = nn.Conv2d(1,128,3)
        self.norm1 = nn.BatchNorm2d(128)
        self.conv2 = nn.Conv2d(128,384,3)
        self.norm2 = nn.BatchNorm2d(384)
        self.conv3 = nn.Conv2d(384,768,3)
        self.norm3 = nn.BatchNorm2d(768)
        self.conv4 = nn.Conv2d(768,2048,3)
        self.conv5 = nn.Conv2d(2048,50,1) # network with 50 outputs
#         self.conv5 = nn.Conv2d(2048,10,1) # network with 10 outputs
        self.pool1 = nn.MaxPool2d((2,4),(2,4))
        self.pool2 = nn.MaxPool2d((2,4),(2,4))
        self.pool3 = nn.MaxPool2d((2,6),(2,6))
        self.pool4 = nn.MaxPool2d((1,7),(1,7))
        self.drop = nn.Dropout2d(.5)
#         self.sig = nn.Sigmoid()
        
    def forward(self,x):
        x = F.relu(self.conv0(x))
#         print(x.shape)
        x = torch.reshape(torch.squeeze(x),(4,1,40,909))
        x = self.pool1(F.relu(self.norm1(self.conv1(x))))
        x = self.pool2(self.drop(F.relu(self.norm2(self.conv2(x)))))
        x = self.pool3(self.drop(F.relu(self.norm3(self.conv3(x)))))
        x = self.pool4(F.relu(self.conv4(x)))
        x = self.conv5(x)
        x = x.view(-1, 50) # network with 50 outputs
#         x = x.view(-1,10) # network with 10 outputs
#         x = self.sig(x)
        return x
    
net = Net()
net.to(device)
# -

# ### Specify file to load pre-saved network parameters IF NETWORK PARAMETERS SAVED
# - specify filename in `filename`
# - following trained networks available
#  - network_10_epoch_10_output_norm_all_ex.pt
#  - network_10_epoch_BNDO_norm_all_ex.pt
#  - network_shuffled.pt
#  - network_10_epoch_1.pt
#  - network_10_epoch_norm_all_ex.pt
#  - network_10_epoch_2.pt
#  - network_10_epoch_norm_binwise.pt

filename = 'network_10_epoch_BNDO_norm_all_ex.pt'
net.load_state_dict(torch.load(os.path.join('./networks',filename)))
net.eval()

# ### creating the training , test and validation dataset

# +
# indices = list(np.arange(len(keys)))
# shuffle(indices)
# train_ind = indices[0:15516]
# val_ind = indices[15517:20689]
# test_ind = indices[20690:]
# print(len(train_ind))
# print(len(val_ind))
# print(len(test_ind))
# with open('train_ind.pickle','wb') as handle:
#     pickle.dump(train_ind,handle)
# with open('val_ind.pickle','wb') as handle:
#     pickle.dump(val_ind,handle)
# with open('test_ind.pickle','wb') as handle:
#     pickle.dump(test_ind,handle)
# -

# - available datasets:
#     - combined_dict_norm_all_examples.pickle
#     - combined_dict_norm_binwise.pickle
#     - combined_dict_norm_per.pickle
#     - combined_dict_updated.pickle - network_10_epoch_1
#     - combined_dict_norm_all_examples_newSpecs.pkl

with open('../database/train_ind.pickle','rb') as handle:
    train_ind = pickle.load(handle)
with open('../database/val_ind.pickle','rb') as handle:
    val_ind = pickle.load(handle)
with open('../database/test_ind.pickle','rb') as handle:
    test_ind = pickle.load(handle)
with open('../database/combined_dict_norm_all_examples.pickle','rb') as handle:
    combined_dict = pickle.load(handle)
with open('../database/sorted_tags.pickle', 'rb') as handle:
    sorted_stats = pickle.load(handle)

# # TRAINING

# ### loading the stored training, test and validation data

# TEST TO CHECK CONTENTS OF DICTIONARY
plt.imshow(combined_dict['2']['mel_spectrum'][0],aspect=20)
plt.show()

# ### Calculating weights for weighted inary Cross Entropy Loss

pos_weight = []
for i in range(50): # make 50 for 50 output
    pos_weight.append(sorted_stats[0][1]/sorted_stats[i][1])
pos_weight = np.array(pos_weight).reshape(1,50) # make 50 for 50 output
print(pos_weight)
pos_weight = torch.from_numpy(pos_weight)
pos_weight = pos_weight.float()
print(type(pos_weight))

# ### Loss function and Optimization functions

# criterion = nn.CrossEntropyLoss()
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight).cuda()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

keys = list(combined_dict.keys())
print(len(keys))

# - comment the line `inputs = inputs.to(device)` and `labels = labels.to(device)` to run on CPU

# +
# remember to change the name of file before execution
batch_size=4
num_channels=1
start_time = time.time()
loss_hist=[]
for epoch in range(10):  # loop over the dataset multiple times
    running_loss = 0.0
    loss_epoch = []
    #creating a mini batch
    for i in range(0,len(train_ind),4):
        shp = combined_dict[keys[train_ind[i]]]['mel_spectrum'][0].shape
        lab_shp = combined_dict[keys[train_ind[i]]]['output'].shape[0] # outputs 50 labels
#         lab_shp=10 # outputs 10 labels
        inputs = np.zeros((batch_size,num_channels,shp[0],shp[1]+1))
        labels = np.zeros((batch_size,lab_shp))
        for j in range(4):
            inputs[j,0,:,0:-1] = combined_dict[keys[train_ind[i+j]]]['mel_spectrum'][0]
            labels[j,:] = combined_dict[keys[train_ind[i+j]]]['output'] # remove last indexing if output is 50 labels
#             labels[j,:] = labels[j,:]*np.arange(50) # was done for crossentropyloss
        inputs = torch.from_numpy(inputs)
        inputs = inputs.float()
        labels = torch.from_numpy(labels)
        labels = labels.float()
        inputs = inputs.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs,labels)
        loss.backward()
        optimizer.step()
        # print statistics
        running_loss += loss.item()
#         print(i)
        if i % 20 == 0:   # every 50 mini-batches 
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 5))
            loss_epoch.append(running_loss/5)
            running_loss = 0.0
    loss_hist.append(loss_epoch)
    torch.save(net.state_dict(),'./net_ep/network_'+str(epoch)+'.pt')
end_time = time.time()
print('Finished Training')
print('Exe. Time:',end_time-start_time)


# -

# ### to save the network parameters onto disk

# change name before saving the network
torch.save(net.state_dict(),'./networks/network_10_epoch_new_model.pt')

# # CODE TO VALIDATION SET PERFORMANCE

keys = list(combined_dict.keys())
print(len(keys))

# - comment line `inputs = inputs.to(device)` to run on CPU

# +
batch_size=4
num_channels=1
out_all = np.zeros((1,50))# change for 50
labels_all = np.zeros((1,50)) # change for 50
for i in range(0,len(val_ind),4):
    shp = combined_dict[keys[val_ind[i]]]['mel_spectrum'][0].shape
    lab_shp = combined_dict[keys[val_ind[i]]]['output'].shape[0]
#     lab_shp = 10 # uncomment for 10 output
    inputs = np.zeros((batch_size,num_channels,shp[0],shp[1]+1))
    labels = np.zeros((batch_size,lab_shp))
    for j in range(4):
        inputs[j,0,:,0:-1] = combined_dict[keys[val_ind[i+j]]]['mel_spectrum'][0]
        labels[j,:] = combined_dict[keys[val_ind[i+j]]]['output'] # remove last indexing if output is 50 labels
    labels_all = np.concatenate((labels_all,labels),axis=0)
    inputs = torch.from_numpy(inputs)
    inputs = inputs.float()
    inputs = inputs.to(device)
    outputs = net(inputs)
    temp = outputs.cpu().detach().numpy()
    out_all = np.concatenate((out_all,temp),axis=0)
    if i%100 == 0:
        print(i)    
    
print('Finished')


# -

labels_all = labels_all[1:,:]
out_all = out_all[1:,:]
print(labels_all.shape)
print(out_all.shape)
out_all_tensor = torch.from_numpy(out_all)
out_all_prob = torch.sigmoid(out_all_tensor).numpy() # for auc_roc
out_all_bin_tensor = torch.round(torch.sigmoid(out_all_tensor))
# out_all_bin = np.where(out_all_prob<.3,0,1) # to change the threshold of rounding
out_all_bin = out_all_bin_tensor.numpy() # for other metrics
print(out_all_bin.shape)

# ### calculating per tag metrics

metrics_dict={}
for i in range(50): # change for 50
#     print(i)
    precision,recall,fscore,support = metrics.precision_recall_fscore_support(labels_all[:,i],out_all_bin[:,i])
    metrics_dict[sorted_stats[i][0]] = {}
    metrics_dict[sorted_stats[i][0]]['precision'] = precision
    metrics_dict[sorted_stats[i][0]]['recall'] = recall
    metrics_dict[sorted_stats[i][0]]['fscore'] = fscore
    metrics_dict[sorted_stats[i][0]]['support'] = support
for key in metrics_dict.keys():
    print(key,':',metrics_dict[key])
    print('\n')


# ### caluclating the metrics(precision,recall,fscore) using all tags at once

precision,recall,fscore,support = metrics.precision_recall_fscore_support(labels_all,out_all_bin)

# ### calculating the AUC-ROC curve

auc_roc = metrics.roc_auc_score(labels_all,out_all_prob)
print(auc_roc)

# # CODE FOR TEST SET

batch_size=4
num_channels=1
out_all = np.zeros((1,50))# change for 50
labels_all = np.zeros((1,50)) # change for 50
auc_all = []
for j in range(10):
    print(j)
    filename = 'network_'+str(j)+'.pt'
    net.load_state_dict(torch.load(os.path.join('./net_ep',filename)))
    net.eval()
    for i in range(0,len(test_ind)-4,4):
        shp = combined_dict[keys[test_ind[i]]]['mel_spectrum'][0].shape
        lab_shp = combined_dict[keys[test_ind[i]]]['output'].shape[0]
    #     lab_shp = 10 # uncomment for 10 output
        inputs = np.zeros((batch_size,num_channels,shp[0],shp[1]+1))
        labels = np.zeros((batch_size,lab_shp))
        for j in range(4):
            inputs[j,0,:,0:-1] = combined_dict[keys[test_ind[i+j]]]['mel_spectrum'][0]
            labels[j,:] = combined_dict[keys[test_ind[i+j]]]['output'] # remove last indexing if output is 50 labels
        labels_all = np.concatenate((labels_all,labels),axis=0)
        inputs = torch.from_numpy(inputs)
        inputs = inputs.float()
        inputs = inputs.to(device)
        outputs = net(inputs)
        temp = outputs.cpu().detach().numpy()
        out_all = np.concatenate((out_all,temp),axis=0)
        if i%100 == 0:
            print(i)
    labels_all = labels_all[1:,:]
    out_all = out_all[1:,:]
    print(labels_all.shape)
    print(out_all.shape)
    out_all_tensor = torch.from_numpy(out_all)
    out_all_prob = torch.sigmoid(out_all_tensor).numpy() # for auc_roc
    out_all_bin_tensor = torch.round(torch.sigmoid(out_all_tensor))
    # out_all_bin = np.where(out_all_prob<.3,0,1) # to change the threshold of rounding
    out_all_bin = out_all_bin_tensor.numpy() # for other metrics
    print(out_all_bin.shape)
    auc_roc = metrics.roc_auc_score(labels_all,out_all_prob)
    auc_all.append(auc_roc)
print('Finished')
    

plt.figure()
plt.plot(auc_all)
plt.xlabel('epochs')
plt.ylabel('auc score')

labels_all = labels_all[1:,:]
out_all = out_all[1:,:]
print(labels_all.shape)
print(out_all.shape)
out_all_tensor = torch.from_numpy(out_all)
out_all_prob = torch.sigmoid(out_all_tensor).numpy() # for auc_roc
out_all_bin_tensor = torch.round(torch.sigmoid(out_all_tensor))
# out_all_bin = np.where(out_all_prob<.3,0,1) # to change the threshold of rounding
out_all_bin = out_all_bin_tensor.numpy() # for other metrics
print(out_all_bin.shape)

labels_all = labels_all[1:,:]
out_all = out_all[1:,:]
print(labels_all.shape)
print(out_all.shape)
out_all_tensor = torch.from_numpy(out_all)
out_all_prob = torch.sigmoid(out_all_tensor).numpy() # for auc_roc
out_all_bin_tensor = torch.round(torch.sigmoid(out_all_tensor))
# out_all_bin = np.where(out_all_prob<.3,0,1) # to change the threshold of rounding
out_all_bin = out_all_bin_tensor.numpy() # for other metrics
print(out_all_bin.shape)

auc_roc = metrics.roc_auc_score(labels_all,out_all_prob)
print(auc_roc)

# - storing all metrics in the dictionary

metrics_dict['all_precesion'] = precision
metrics_dict['all_recall'] = recall
metrics_dict['all_fscore'] = fscore
metrics_dict['all_support'] = support
metrics_dict['auc_roc'] = auc_roc
metrics_dict['loss_hist'] = loss_hist

# saving the metric to disk
with open('./metrics/metrics_10_epoch_BNDO_norm_all_ex_test.pickle','wb') as handle:
    pickle.dump(metrics_dict,handle)

StarcoderdataJupyterScriptsDedupFiltered