Update inference_2.py

inference_2.py

@@ -1,22 +1,32 @@
+import os
 import cv2
 import onnx
 import torch
 import argparse
 import numpy as np
+import torch.nn as nn
+from models.TMC import ETMC
 from models import image
 
-import warnings
 from onnx2pytorch import ConvertModel
 
-warnings.filterwarnings("ignore", message="The given NumPy array is not writable")
-with warnings.catch_warnings():
-    warnings.filterwarnings("ignore", message="The given NumPy array is not writable")
-    onnx_model = onnx.load('models/efficientnet.onnx')
-    pytorch_model = ConvertModel(onnx_model)
+onnx_model = onnx.load('checkpoints/efficientnet.onnx')
+pytorch_model = ConvertModel(onnx_model)
+
 torch.manual_seed(42)
 
 
-audio_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}
+audio_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
+}
 
 
 def get_args(parser):
@@ -55,14 +65,14 @@ def get_args(parser):
 
 def load_img_modality_model(args):
     rgb_encoder = pytorch_model
-    ckpt = torch.load('models/model.pth', map_location = torch.device('cpu'))
+    ckpt = torch.load('checkpoints/model.pth', map_location = torch.device('cpu'))
     rgb_encoder.load_state_dict(ckpt['rgb_encoder'], strict = True)
     rgb_encoder.eval()
     return rgb_encoder
 
 def load_spec_modality_model(args):
     spec_encoder = image.RawNet(args)
-    ckpt = torch.load('models/model.pth', map_location = torch.device('cpu'))
+    ckpt = torch.load('checkpoints/model.pth', map_location = torch.device('cpu'))
     spec_encoder.load_state_dict(ckpt['spec_encoder'], strict = True)
     spec_encoder.eval()
     return spec_encoder
@@ -71,9 +81,12 @@ parser = argparse.ArgumentParser(description="Inference models")
 get_args(parser)
 args, remaining_args = parser.parse_known_args()
 assert remaining_args == [], remaining_args
+
 spec_model = load_spec_modality_model(args)
+
 img_model = load_img_modality_model(args)
 
+
 def preprocess_img(face):
     face = face / 255
     face = cv2.resize(face, (256, 256))
@@ -90,12 +103,15 @@ def df_spec_pred(input_audio):
     spec_grads = spec_model.forward(audio)
     spec_grads_inv = np.exp(spec_grads.cpu().detach().numpy().squeeze())
     max_value = np.argmax(spec_grads_inv)
+
     if max_value > 0.5:
         preds = round(100 - (max_value*100), 3)
         text2 = f"The audio is REAL."
+
     else:
         preds = round(max_value*100, 3)
         text2 = f"The audio is FAKE."
+
     return text2
 
 def df_img_pred(input_image):
@@ -104,25 +120,34 @@ def df_img_pred(input_image):
     img_grads = img_model.forward(face)
     img_grads = img_grads.cpu().detach().numpy()
     img_grads_np = np.squeeze(img_grads)
+
     if img_grads_np[0] > 0.5:
         preds = round(img_grads_np[0] * 100, 3)
         text2 = f"The image is REAL. \nConfidence score is: {preds}"
+
     else:
         preds = round(img_grads_np[1] * 100, 3)
         text2 = f"The image is FAKE. \nConfidence score is: {preds}"
+
     return text2
 
+
 def preprocess_video(input_video, n_frames = 3):
     v_cap = cv2.VideoCapture(input_video)
     v_len = int(v_cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    # Pick 'n_frames' evenly spaced frames to sample
     if n_frames is None:
         sample = np.arange(0, v_len)
     else:
         sample = np.linspace(0, v_len - 1, n_frames).astype(int)
+
+    #Loop through frames.
     frames = []
     for j in range(v_len):
         success = v_cap.grab()
         if j in sample:
+            # Load frame
             success, frame = v_cap.retrieve()
             if not success:
                 continue
@@ -132,22 +157,28 @@ def preprocess_video(input_video, n_frames = 3):
     v_cap.release()
     return frames
 
+
 def df_video_pred(input_video):
     video_frames = preprocess_video(input_video)
     real_faces_list = []
     fake_faces_list = []
+
     for face in video_frames:
         img_grads = img_model.forward(face)
         img_grads = img_grads.cpu().detach().numpy()
         img_grads_np = np.squeeze(img_grads)
         real_faces_list.append(img_grads_np[0])
         fake_faces_list.append(img_grads_np[1])
+
     real_faces_mean = np.mean(real_faces_list)
     fake_faces_mean = np.mean(fake_faces_list)
+
     if real_faces_mean > 0.5:
         preds = round(real_faces_mean * 100, 3)
         text2 = f"The video is REAL. \nConfidence score is: {preds}%"
+
     else:
         preds = round(fake_faces_mean * 100, 3)
         text2 = f"The video is FAKE. \nConfidence score is: {preds}%"
+
     return text2