add

app.py

@@ -5,7 +5,8 @@ import sounddevice as sd
 import numpy as np
 import pandas as pd
 import torch
-# import torchaudio
+import torch.nn as nn
+import torchaudio
 import wave
 import io
 from scipy.io import wavfile
@@ -14,14 +15,38 @@ import time
 import os
 import atexit
 import librosa
+import torchaudio.functional as F
+import torchaudio.transforms as T
 
 # MODEL LOADING and INITIALISATION
-
-def load_model():
-    model = torch.jit.load("snorenetv1_small.ptl")
-    model.eval()
-    return model
-model = load_model()
+n_fft = 1024
+win_length = None
+hop_length = 32
+
+# Input tensor shape was ([1,16000])
+class SnoreNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.transform = torchaudio.transforms.Spectrogram(n_fft = n_fft,win_length = win_length,hop_length = hop_length,center = True ,pad_mode = "reflect",power = 2.0,)
+        self.fc1 = nn.Linear(257013, 512)
+        self.act1 = nn.Tanh()
+        self.fc2 = nn.Linear(512, 2)
+        self.logs1 = nn.LogSoftmax(dim=1)
+    
+    def forward(self, raw_audio_tensor):
+        # print(raw_audio_tensor.shape)
+        spectrogram = self.transform(raw_audio_tensor)
+        # print(spectrogram.shape)
+        spectrogram = spectrogram.reshape(spectrogram.size(0), -1)
+        # print(spectrogram.shape)
+        output = self.act1(self.fc1(spectrogram))
+        output = self.fc2(output)
+        output = torch.abs(self.logs1(output))
+        return output
+    
+model = SnoreNet()
+model.load_state_dict(torch.load('snoreNetv1.pt'))
+model.eval()
 
 
 # Audio parameters
@@ -31,7 +56,8 @@ def process_data(waveform_chunks):
     for chunk in waveform_chunks:
         input_tensor = torch.tensor(chunk).unsqueeze(0).to(torch.float32)
         # st.write(input_tensor[0][98])
-        result = model(input_tensor)
+        with torch.no_grad():
+            result = model(input_tensor)
         # st.write(result)
         if np.abs(result[0][0]) > np.abs(result[0][1]):
             other += 1