ADD: demo

app.py

@@ -0,0 +1,105 @@
+import gradio as gr
+import cv2
+from PIL import Image
+import torch
+import numpy as np
+
+from transformers import AutoImageProcessor, AutoProcessor, AutoModel, CLIPVisionModel
+from detection import detect_image, detect_video
+from model import LinearClassifier
+
+
+def load_model(detection_type):
+
+    device = torch.device("cpu")
+
+    processor = AutoProcessor.from_pretrained("openai/clip-vit-large-patch14")
+    clip_model = CLIPVisionModel.from_pretrained("openai/clip-vit-large-patch14", output_attentions=True)
+
+    model_path = f"pretrained_models/{detection_type}/clip_weights.pth"
+    checkpoint = torch.load(model_path, map_location="cpu")
+    input_dim = checkpoint["linear.weight"].shape[1]
+    
+    detection_model = LinearClassifier(input_dim)
+    detection_model.load_state_dict(checkpoint)
+    detection_model = detection_model.to(device)
+
+    return processor, clip_model, detection_model
+
+def process_image(image, detection_type):
+    processor, clip_model, detection_model = load_model(detection_type)
+    
+    results = detect_image(image, processor, clip_model, detection_model)
+
+    pred_score = results["pred_score"]
+    attn_map = results["attn_map"]
+
+    return pred_score, attn_map
+
+def process_video(video, detection_type):
+    processor, clip_model, detection_model = load_model(detection_type)
+
+    cap = cv2.VideoCapture(video)
+    frames = []
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        pil_image = Image.fromarray(frame)
+        frames.append(pil_image)
+    cap.release()
+
+    results = detect_video(frames, processor, clip_model, detection_model)
+
+    pred_score = results["pred_score"]
+    attn_map = results["attn_map"]
+
+    return pred_score, attn_map
+
+def change_input(input_type):
+    if input_type == "Image":
+        return gr.update(visible=True), gr.update(visible=False)
+    elif input_type == "Video":
+        return gr.update(visible=False), gr.update(visible=True)
+    else:
+        return None
+
+
+def process_input(input_type, model_type, image, video):
+    detection_type = "facial" if model_type == "Facial" else "general"
+
+    if input_type == "Image" and image is not None:
+        return process_image(image, detection_type)
+    elif input_type == "Video" and video is not None:
+        return process_video(video, detection_type)
+    else:
+        return None, None
+
+
+with gr.Blocks() as demo:
+  
+    gr.Markdown("## Deepfake Detection : Facial / General")
+  
+    input_type = gr.Radio(["Image", "Video"], label="Choose Input Type", value="Image")
+
+    model_type = gr.Radio(["Facial", "General"], label="Choose Model Type", value="General")
+  
+    image_input = gr.Image(type="pil", label="Upload Image", visible=True)
+    video_input = gr.Video(label="Upload Video", visible=False)
+  
+    process_button = gr.Button("Run Model")
+
+    pred_score_output = gr.Textbox(label="Prediction Score")
+    attn_map_output = gr.Image(type="pil", label="Attention Map")
+  
+    input_type.change(fn=change_input, inputs=[input_type], outputs=[image_input, video_input])
+  
+    process_button.click(
+        fn=process_input, 
+        inputs=[input_type, model_type, image_input, video_input], 
+        outputs=[pred_score_output, attn_map_output]
+    )
+
+if __name__ == "__main__":
+    demo.launch()

detection.py

@@ -0,0 +1,85 @@
+import torch
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+from PIL import Image
+
+
+def vis_attn(image, patch_attention_map, alpha=0.5, vis_option="none"):
+
+    image = np.array(image)
+    H, W, _ = image.shape
+
+    seq_len = patch_attention_map.shape[0]
+    grid_size = int(seq_len ** 0.5)
+
+    patch_attention_map = patch_attention_map.reshape(grid_size, grid_size)
+    patch_attention_map = cv2.resize(patch_attention_map.cpu().detach().numpy(), (W, H), interpolation=cv2.INTER_CUBIC)
+
+    patch_attention_map = (patch_attention_map - patch_attention_map.min()) / (patch_attention_map.max() - patch_attention_map.min())
+    patch_attention_map = np.uint8(255 * patch_attention_map)
+
+    heatmap = cv2.applyColorMap(patch_attention_map, cv2.COLORMAP_JET)
+
+    blended_image = cv2.addWeighted(image, 1 - alpha, heatmap, alpha, 0)
+    blended_image = cv2.cvtColor(blended_image, cv2.COLOR_RGB2BGR)
+    blended_image = Image.fromarray(blended_image)
+
+    return blended_image
+
+
+def detect_image(image, processor, clip_model, detection_model):
+
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = clip_model(**inputs)
+    
+    last_hidden_states = outputs.last_hidden_state[:, 0, :]
+    
+    pred_score = float(detection_model(last_hidden_states)[0][0].cpu().detach().numpy())
+    assert 0 <= pred_score <= 1
+    
+    for layer_idx in range(len(outputs.attentions)):
+        attn_map = outputs.attentions[layer_idx]
+        if layer_idx == 0:
+            last_layer_attn = attn_map
+        else:
+            if layer_idx < 6:
+                last_layer_attn += attn_map
+    
+    head_mean_attn = last_layer_attn.mean(dim=1)[0]
+    cls_attention_map = head_mean_attn[0, 1:]
+
+    blended_image = vis_attn(image, cls_attention_map)
+
+    results = {
+        "pred_score": pred_score,
+        "attn_map": blended_image,
+    }
+
+    return results
+
+def detect_video(frames, processor, clip_model, detection_model):
+    image = frames[0]
+
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = clip_model(**inputs)
+    
+    last_hidden_states = outputs.last_hidden_state[:, 0, :]
+    
+    pred_score = float(detection_model(last_hidden_states)[0][0].cpu().detach().numpy())
+    assert 0 <= pred_score <= 1
+
+    attention_maps = outputs.attentions[-1].cpu().detach().numpy()
+    cls_attention_map = attention_maps[:, :, 0, :]
+    cls_attention_map = cls_attention_map.mean(axis=0)
+
+    blended_image = vis_attn(image, cls_attention_map)
+
+    results = {
+        "pred_score": pred_score,
+        "attn_map": blended_image,
+    }
+
+    return results

model.py

@@ -0,0 +1,12 @@
+import torch
+
+class LinearClassifier(torch.nn.Module):
+    def __init__(self, input_dim):
+        super(LinearClassifier, self).__init__()
+        self.linear = torch.nn.Linear(input_dim, 1)
+        self.sigmoid = torch.nn.Sigmoid()
+
+    def forward(self, x):
+        x = self.linear(x)
+        x = self.sigmoid(x)
+        return x

pretrained_models/facial/clip_weights.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59d99962ba4c7697416755ec815dc355d9694f957e4a76806e891a669bb33c5b
+size 5686

pretrained_models/facial/dino_weights.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7bdec7ff5d97ba4101085352756a546c15829d9e5910934e66fd7cc13eb5458f
+size 5686

pretrained_models/facial/mobileclip_weights.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e94c2912a932b906dc0b2cd062b1c591948656627e48ce187ab6a3dd508a8a4
+size 3674

pretrained_models/general/clip_weights.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:189df51aab3791bea65305e3c2807341fea7bec4cb2019cc8e58d6b217c59deb
+size 5686

pretrained_models/general/dino_weights.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:43059530466fce466a0e242a1523936a9bb2f60782915ad42837f974dbbcaad1
+size 5674

pretrained_models/general/mobileclip_weights.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c09f581e3fb36c74444b74f172f8ae5a601e53071b61f855f626ac37fd687875
+size 3674