Update pipeline.py

pipeline.py

@@ -165,5 +165,48 @@ def deepfakes_image_predict(input_image):
     else:
         text2 = "The image is FAKE. \n Deepfakes Confidence: " + str(round(fake*100, 3)) + "%"
     return text2
+    
+def load_audio_model():
+    d_args = {
+  "nb_samp": 64600,
+  "first_conv": 1024,
+  "in_channels": 1,
+  "filts": [20, [20, 20], [20, 128], [128, 128]],
+  "blocks": [2, 4],
+  "nb_fc_node": 1024,
+  "gru_node": 1024,
+  "nb_gru_layer": 3,
+  "nb_classes": 2}
+    
+    model = RawNet(d_args = d_args, device='cpu')
+
+    #Load ckpt.
+    model_dict = model.state_dict()
+    ckpt = torch.load('RawNet2.pth', map_location=torch.device('cpu'))
+    model.load_state_dict(ckpt, model_dict)
+    return model
+
+audio_label_map = {
+    0: "Real audio",
+    1: "Fake audio"
+}
+
+def deepfakes_audio_predict(input_audio):
+    #Perform inference on audio.
+    x, sr = input_audio
+    x_pt = torch.Tensor(x)
+    x_pt = torch.unsqueeze(x_pt, dim = 0)
+
+    #Load model.
+    model = load_audio_model()
+
+    #Perform inference.
+    grads = model(x_pt)
+
+    #Get the argmax.
+    grads_np = grads.detach().numpy()
+    result = np.argmax(grads_np)
+
+    return audio_label_map[result]