aneela-pervez commited on
Commit
36bb346
·
verified ·
1 Parent(s): 12d0475

Create fake_image.py

Browse files
Files changed (1) hide show
  1. fake_image.py +100 -0
fake_image.py ADDED
@@ -0,0 +1,100 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # -*- coding: utf-8 -*-
2
+ import torch
3
+ import torch.nn as nn
4
+ import torchvision.transforms as transforms
5
+ import torchvision.models as models
6
+ from PIL import Image
7
+ import os
8
+ import numpy as np
9
+ from huggingface_hub import hf_hub_download
10
+
11
+ # --- 🚨 MASTER FIX FOR PYTORCH 2.6 SECURITY ---
12
+ import torch.serialization
13
+ try:
14
+ torch.serialization.add_safe_globals([np.core.multiarray.scalar])
15
+ except:
16
+ pass
17
+
18
+ _original_load = torch.load
19
+ def _patched_load(*args, **kwargs):
20
+ kwargs['weights_only'] = False
21
+ return _original_load(*args, **kwargs)
22
+ torch.load = _patched_load
23
+ # ----------------------------------------------
24
+
25
+ # 1. Device Configuration
26
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
27
+
28
+ # 2. Model Architecture (SOTA Hybrid Network)
29
+ class Deepfake_Hybrid_Network(nn.Module):
30
+ def __init__(self, num_classes=2):
31
+ super(Deepfake_Hybrid_Network, self).__init__()
32
+ resnet = models.resnet18(weights=None)
33
+ self.cnn_extractor = nn.Sequential(*list(resnet.children())[:-2])
34
+
35
+ encoder_layer = nn.TransformerEncoderLayer(
36
+ d_model=512, nhead=8, dim_feedforward=1024, dropout=0.3, batch_first=True
37
+ )
38
+ self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=2)
39
+
40
+ self.classifier = nn.Sequential(
41
+ nn.Linear(512, 128),
42
+ nn.BatchNorm1d(128),
43
+ nn.ReLU(),
44
+ nn.Dropout(0.4),
45
+ nn.Linear(128, num_classes)
46
+ )
47
+
48
+ def forward(self, x):
49
+ x = self.cnn_extractor(x)
50
+ b, c, h, w = x.shape
51
+ x = x.view(b, c, h * w).permute(0, 2, 1)
52
+ x = self.transformer_encoder(x)
53
+ x = x.mean(dim=1)
54
+ return self.classifier(x)
55
+
56
+ # 3. Load Model from Hugging Face Hub
57
+ print("Loading Image Hybrid SOTA Model...")
58
+ model = Deepfake_Hybrid_Network(num_classes=2).to(device)
59
+
60
+ try:
61
+ # Model HF repo se download ho raha hy
62
+ model_path = hf_hub_download(repo_id="aneela-pervez/My-Deepfake-Models", filename="Deepfake_Hybrid_SOTA.pth")
63
+ state_dict = torch.load(model_path, map_location=device)
64
+
65
+ if any(k.startswith('module.') for k in state_dict.keys()):
66
+ from collections import OrderedDict
67
+ new_state_dict = OrderedDict()
68
+ for k, v in state_dict.items():
69
+ new_state_dict[k.replace('module.', '')] = v
70
+ model.load_state_dict(new_state_dict)
71
+ else:
72
+ model.load_state_dict(state_dict)
73
+ model.eval()
74
+ print("✅ Image SOTA Model loaded successfully!")
75
+ except Exception as e:
76
+ print(f"⚠️ Error loading image model weights: {e}")
77
+
78
+ # 4. Image Preprocessing
79
+ test_transform = transforms.Compose([
80
+ transforms.Resize((224, 224)),
81
+ transforms.ToTensor(),
82
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
83
+ ])
84
+
85
+ # 5. Prediction Function for Master Pipeline
86
+ def predict_image(img):
87
+ if img is None:
88
+ return {"Error": 1.0}
89
+
90
+ img_tensor = test_transform(img).unsqueeze(0).to(device)
91
+
92
+ with torch.no_grad():
93
+ outputs = model(img_tensor)
94
+ probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
95
+
96
+ results = {
97
+ "Fake (Deepfake)": float(probabilities[0]),
98
+ "Real Image": float(probabilities[1])
99
+ }
100
+ return results