Update inference_2.py

inference_2.py

@@ -1,216 +1,216 @@
-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
-
-from onnx2pytorch import ConvertModel
-
-onnx_model = onnx.load('checkpoints/efficientnet.onnx')
-pytorch_model = ConvertModel(onnx_model)
-
-#Set random seed for reproducibility.
-torch.manual_seed(42)
-
-
-# Define the audio_args dictionary
-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):
-    parser.add_argument("--batch_size", type=int, default=8)
-    parser.add_argument("--data_dir", type=str, default="datasets/train/fakeavceleb*")
-    parser.add_argument("--LOAD_SIZE", type=int, default=256)
-    parser.add_argument("--FINE_SIZE", type=int, default=224)
-    parser.add_argument("--dropout", type=float, default=0.2)
-    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
-    parser.add_argument("--hidden", nargs="*", type=int, default=[])
-    parser.add_argument("--hidden_sz", type=int, default=768)
-    parser.add_argument("--img_embed_pool_type", type=str, default="avg", choices=["max", "avg"])
-    parser.add_argument("--img_hidden_sz", type=int, default=1024)
-    parser.add_argument("--include_bn", type=int, default=True)
-    parser.add_argument("--lr", type=float, default=1e-4)
-    parser.add_argument("--lr_factor", type=float, default=0.3)
-    parser.add_argument("--lr_patience", type=int, default=10)
-    parser.add_argument("--max_epochs", type=int, default=500)
-    parser.add_argument("--n_workers", type=int, default=12)
-    parser.add_argument("--name", type=str, default="MMDF")
-    parser.add_argument("--num_image_embeds", type=int, default=1)
-    parser.add_argument("--patience", type=int, default=20)
-    parser.add_argument("--savedir", type=str, default="./savepath/")
-    parser.add_argument("--seed", type=int, default=1)
-    parser.add_argument("--n_classes", type=int, default=2)
-    parser.add_argument("--annealing_epoch", type=int, default=10)
-    parser.add_argument("--device", type=str, default='cpu')
-    parser.add_argument("--pretrained_image_encoder", type=bool, default = False)
-    parser.add_argument("--freeze_image_encoder", type=bool, default = False)
-    parser.add_argument("--pretrained_audio_encoder", type = bool, default=False)
-    parser.add_argument("--freeze_audio_encoder", type = bool, default = False)
-    parser.add_argument("--augment_dataset", type = bool, default = True)
-
-    for key, value in audio_args.items():
-        parser.add_argument(f"--{key}", type=type(value), default=value)
-
-def model_summary(args):
-    '''Prints the model summary.'''
-    model = ETMC(args)
-
-    for name, layer in model.named_modules():
-        print(name, layer)
-
-def load_multimodal_model(args):
-    '''Load multimodal model'''
-    model = ETMC(args)
-    ckpt = torch.load('checkpoints/model.pth', map_location = torch.device('cpu'))
-    model.load_state_dict(ckpt, strict = True)
-    model.eval()
-    return model
-
-def load_img_modality_model(args):
-    '''Loads image modality model.'''
-    rgb_encoder = pytorch_model
-
-    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('checkpoints/model.pth', map_location = torch.device('cpu'))
-    spec_encoder.load_state_dict(ckpt['spec_encoder'], strict = True)
-    spec_encoder.eval()
-    return spec_encoder
-
-
-#Load models.
-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))
-    # face = face.transpose(2, 0, 1) #(W, H, C) -> (C, W, H)
-    face_pt = torch.unsqueeze(torch.Tensor(face), dim = 0) 
-    return face_pt
-
-def preprocess_audio(audio_file):
-    audio_pt = torch.unsqueeze(torch.Tensor(audio_file), dim = 0)
-    return audio_pt
-
-def deepfakes_spec_predict(input_audio):
-    x, _ = input_audio
-    audio = preprocess_audio(x)
-    spec_grads = spec_model.forward(audio)
-    spec_grads_inv = np.exp(spec_grads.cpu().detach().numpy().squeeze())
-
-    # multimodal_grads = multimodal.spec_depth[0].forward(spec_grads)
-
-    # out = nn.Softmax()(multimodal_grads)
-    # max = torch.argmax(out, dim = -1) #Index of the max value in the tensor.
-    # max_value = out[max] #Actual value of the tensor.
-    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 deepfakes_image_predict(input_image):
-    face = preprocess_img(input_image)
-    print(f"Face shape is: {face.shape}")
-    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
-            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            frame = preprocess_img(frame)
-            frames.append(frame)
-    v_cap.release()
-    return frames
-
-
-def deepfakes_video_predict(input_video):
-    '''Perform inference on a video.'''
-    video_frames = preprocess_video(input_video)
-    real_faces_list = []
-    fake_faces_list = []
-
-    for face in video_frames:
-        # face = preprocess_img(face)
-
-        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
-
+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
+
+from onnx2pytorch import ConvertModel
+
+onnx_model = onnx.load('checkpoints/efficientnet.onnx')
+pytorch_model = ConvertModel(onnx_model)
+
+#Set random seed for reproducibility.
+torch.manual_seed(42)
+
+
+# Define the audio_args dictionary
+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):
+    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--data_dir", type=str, default="datasets/train/fakeavceleb*")
+    parser.add_argument("--LOAD_SIZE", type=int, default=256)
+    parser.add_argument("--FINE_SIZE", type=int, default=224)
+    parser.add_argument("--dropout", type=float, default=0.2)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
+    parser.add_argument("--hidden", nargs="*", type=int, default=[])
+    parser.add_argument("--hidden_sz", type=int, default=768)
+    parser.add_argument("--img_embed_pool_type", type=str, default="avg", choices=["max", "avg"])
+    parser.add_argument("--img_hidden_sz", type=int, default=1024)
+    parser.add_argument("--include_bn", type=int, default=True)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--lr_factor", type=float, default=0.3)
+    parser.add_argument("--lr_patience", type=int, default=10)
+    parser.add_argument("--max_epochs", type=int, default=500)
+    parser.add_argument("--n_workers", type=int, default=12)
+    parser.add_argument("--name", type=str, default="MMDF")
+    parser.add_argument("--num_image_embeds", type=int, default=1)
+    parser.add_argument("--patience", type=int, default=20)
+    parser.add_argument("--savedir", type=str, default="./savepath/")
+    parser.add_argument("--seed", type=int, default=1)
+    parser.add_argument("--n_classes", type=int, default=2)
+    parser.add_argument("--annealing_epoch", type=int, default=10)
+    parser.add_argument("--device", type=str, default='cpu')
+    parser.add_argument("--pretrained_image_encoder", type=bool, default = False)
+    parser.add_argument("--freeze_image_encoder", type=bool, default = False)
+    parser.add_argument("--pretrained_audio_encoder", type = bool, default=False)
+    parser.add_argument("--freeze_audio_encoder", type = bool, default = False)
+    parser.add_argument("--augment_dataset", type = bool, default = True)
+
+    for key, value in audio_args.items():
+        parser.add_argument(f"--{key}", type=type(value), default=value)
+
+def model_summary(args):
+    '''Prints the model summary.'''
+    model = ETMC(args)
+
+    for name, layer in model.named_modules():
+        print(name, layer)
+
+def load_multimodal_model(args):
+    '''Load multimodal model'''
+    model = ETMC(args)
+    ckpt = torch.load('checkpoints/model.pth', map_location = torch.device('cpu'))
+    model.load_state_dict(ckpt, strict = True)
+    model.eval()
+    return model
+
+def load_img_modality_model(args):
+    '''Loads image modality model.'''
+    rgb_encoder = pytorch_model
+
+    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('checkpoints/model.pth', map_location = torch.device('cpu'))
+    spec_encoder.load_state_dict(ckpt['spec_encoder'], strict = True)
+    spec_encoder.eval()
+    return spec_encoder
+
+
+#Load models.
+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))
+    # face = face.transpose(2, 0, 1) #(W, H, C) -> (C, W, H)
+    face_pt = torch.unsqueeze(torch.Tensor(face), dim = 0) 
+    return face_pt
+
+def preprocess_audio(audio_file):
+    audio_pt = torch.unsqueeze(torch.Tensor(audio_file), dim = 0)
+    return audio_pt
+
+def deepfakes_spec_predict(input_audio):
+    x, _ = input_audio
+    audio = preprocess_audio(x)
+    spec_grads = spec_model.forward(audio)
+    spec_grads_inv = np.exp(spec_grads.cpu().detach().numpy().squeeze())
+
+    # multimodal_grads = multimodal.spec_depth[0].forward(spec_grads)
+
+    # out = nn.Softmax()(multimodal_grads)
+    # max = torch.argmax(out, dim = -1) #Index of the max value in the tensor.
+    # max_value = out[max] #Actual value of the tensor.
+    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 deepfakes_image_predict(input_image):
+    face = preprocess_img(input_image)
+    print(f"Face shape is: {face.shape}")
+    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."
+
+    else:
+        preds = round(img_grads_np[1] * 100, 3)
+        text2 = f"The image is FAKE."
+
+    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
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame = preprocess_img(frame)
+            frames.append(frame)
+    v_cap.release()
+    return frames
+
+
+def deepfakes_video_predict(input_video):
+    '''Perform inference on a video.'''
+    video_frames = preprocess_video(input_video)
+    real_faces_list = []
+    fake_faces_list = []
+
+    for face in video_frames:
+        # face = preprocess_img(face)
+
+        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."
+
+    else:
+        preds = round(fake_faces_mean * 100, 3)
+        text2 = f"The video is FAKE."
+
+    return text2
+