Update model and API

cnn.py

@@ -52,7 +52,7 @@ class CNNetwork(nn.Module):
             nn.MaxPool2d(kernel_size=2)
         )
         self.flatten = nn.Flatten()
-        self.linear = nn.Linear(128 * 9 * 31, 3)
+        self.linear = nn.Linear(128 * 9 * 19, 3)
         self.softmax = nn.Softmax(dim=1)
 
     def forward(self, input_data):

server/main.py

@@ -14,7 +14,7 @@ from cnn import CNNetwork
 
 # load model
 model = CNNetwork()
-state_dict = torch.load("../models/aisf/void_20230516_193200.pth")
+state_dict = torch.load("../models/aisf/void_20230517_102846.pth")
 model.load_state_dict(state_dict)
 
 # TODO: update to grabbing labels stored on model
@@ -54,8 +54,11 @@ def model_predict(wav):
     model_input = wav.unsqueeze(0)
     output = model(model_input)
     prediction_index = torch.argmax(output, 1).item()
+    output = output.detach().cpu().numpy()[0]
 
     return {
         "prediction_index": prediction_index,
-        "prediciton": LABELS[prediction_index],
+        "labels": LABELS,
+        "prediction_label": LABELS[prediction_index],
+        "prediction_output": output.tolist(),
     }

server/preprocess.py

@@ -9,6 +9,7 @@ from scipy.io import wavfile
 import wget
 
 DEFAULT_SAMPLE_RATE=48000
+DEFAULT_WAVE_LENGTH=3
 
 def process_from_url(url):
     # download UI audio
@@ -26,7 +27,7 @@ def process_from_url(url):
     return spec
 
 
-def process_from_filename(filename, target_sample_rate=DEFAULT_SAMPLE_RATE, wav_length=5):
+def process_from_filename(filename, target_sample_rate=DEFAULT_SAMPLE_RATE, wav_length=DEFAULT_WAVE_LENGTH):
     wav, sample_rate = torchaudio.load(filename)
 
     wav = process_raw_wav(wav, sample_rate, target_sample_rate, wav_length)
@@ -35,7 +36,7 @@ def process_from_filename(filename, target_sample_rate=DEFAULT_SAMPLE_RATE, wav_
 
     return spec
 
-def process_raw_wav(wav, sample_rate=DEFAULT_SAMPLE_RATE, target_sample_rate=DEFAULT_SAMPLE_RATE, wav_length=5):
+def process_raw_wav(wav, sample_rate=DEFAULT_SAMPLE_RATE, target_sample_rate=DEFAULT_SAMPLE_RATE, wav_length=DEFAULT_WAVE_LENGTH):
     num_samples = wav_length * target_sample_rate
 
     wav = _resample(wav, sample_rate, target_sample_rate)