Spaces:
Sleeping
Sleeping
File size: 5,710 Bytes
11aa20d 823741b 11aa20d 2df033f 11aa20d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """
Vision-KAN + Vision-LSTM hybrid deepfake detector.
This module is a faithful, refactored copy of the inference path in
`notebooks/test_mac.ipynb`. The architecture, weights loading, preprocessing,
and the real/fake decision rule are IDENTICAL to the original notebook — only
reorganized into importable functions so the demo app and the notebook share
one source of truth.
Decision rule (unchanged from the original project):
prob = sigmoid output of the model
prob >= 0.5 -> REAL
prob < 0.5 -> FAKE
"""
import os
import numpy as np
import torch
import torch.nn as nn
from PIL import Image
DEVICE = torch.device("cpu")
# ImageNet normalization — exactly as used during training/inference.
_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3)
_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3)
# ──────────────────────────────────────────────────────────────────────────
# Model definition (unchanged from test_mac.ipynb)
# ──────────────────────────────────────────────────────────────────────────
class CombinedModel(nn.Module):
"""Late-fusion of a Vision-KAN backbone and a Vision-LSTM backbone.
Each backbone emits a scalar logit. The two logits are concatenated and
passed through a small MLP, then a sigmoid produces a probability in [0, 1].
"""
def __init__(self, KAN, vision_lstm):
super(CombinedModel, self).__init__()
self.KAN = KAN
self.vision_lstm = vision_lstm
self.fc1 = nn.Linear(2, 512)
self.fc2 = nn.Linear(512, 1)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
KAN_output = self.KAN(x)
lstm_output = self.vision_lstm(x)
combined = torch.cat((KAN_output, lstm_output), dim=1)
x = self.fc1(combined)
x = self.fc2(x)
x = self.sigmoid(x)
return x
# ──────────────────────────────────────────────────────────────────────────
# Build + load
# ──────────────────────────────────────────────────────────────────────────
def build_model():
"""Reconstruct the exact Vision-KAN + Vision-LSTM hybrid (random weights)."""
from VisionKAN import create_model
KAN = create_model(
model_name="deit_tiny_patch16_224_KAN",
pretrained=False,
hdim_kan=192,
num_classes=1,
drop_rate=0.0,
drop_path_rate=0.05,
img_size=224,
batch_size=32,
)
vision_lstm = torch.hub.load(
"nx-ai/vision-lstm",
"VisionLSTM",
dim=192,
depth=24,
patch_size=16,
input_shape=(3, 224, 224),
output_shape=(1,),
drop_path_rate=0.05,
stride=None,
trust_repo=True,
skip_validation=True, # avoid GitHub API rate limit on shared Space IPs
)
return CombinedModel(KAN, vision_lstm).to(DEVICE)
def load_model(weights_path="as_model_best.pt"):
"""Build the model and load trained weights, returning it in eval mode."""
model = build_model()
if not os.path.exists(weights_path):
raise FileNotFoundError(
f"Trained weights not found at '{weights_path}'. "
f"Download as_model_best.pt (see README) and place it here."
)
state_dict = torch.load(weights_path, map_location=DEVICE)
model.load_state_dict(state_dict)
model.eval()
return model
# ──────────────────────────────────────────────────────────────────────────
# Preprocessing (unchanged from test_mac.ipynb -> preprocess_pil)
# ──────────────────────────────────────────────────────────────────────────
def preprocess_pil(img: Image.Image) -> torch.Tensor:
"""Resize to 224x224, scale to [0,1], ImageNet-normalize, return CHW tensor."""
img = img.convert("RGB").resize((224, 224))
img_np = np.asarray(img, dtype=np.float32) / 255.0
img_np = (img_np - _MEAN) / _STD
img_np = img_np.transpose(2, 0, 1)
return torch.tensor(img_np, dtype=torch.float32)
# ──────────────────────────────────────────────────────────────────────────
# Prediction
# ──────────────────────────────────────────────────────────────────────────
@torch.no_grad()
def predict(model, img: Image.Image):
"""Run the model on a PIL image.
Returns (label, probability) where probability is the raw sigmoid output
and label is decided with the original threshold: prob >= 0.5 -> REAL.
"""
x = preprocess_pil(img).unsqueeze(0).to(DEVICE)
prob = model(x).item()
label = "REAL" if prob >= 0.5 else "FAKE"
return label, prob
|