atharvawarade9807 commited on
Commit
b3a532e
·
verified ·
1 Parent(s): f76a5c7

Upload 11 files

Browse files
Version_2 - Copy/core/__init__.py ADDED
File without changes
Version_2 - Copy/core/__pycache__/__init__.cpython-313.pyc ADDED
Binary file (147 Bytes). View file
 
Version_2 - Copy/core/__pycache__/augmentations.cpython-313.pyc ADDED
Binary file (1.98 kB). View file
 
Version_2 - Copy/core/__pycache__/data_loader.cpython-313.pyc ADDED
Binary file (3.48 kB). View file
 
Version_2 - Copy/core/__pycache__/sam.cpython-313.pyc ADDED
Binary file (3.64 kB). View file
 
Version_2 - Copy/core/augmentations.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from torchvision.transforms import v2
3
+
4
+ def get_train_transforms(image_size=224):
5
+ return v2.Compose([
6
+ v2.RandomResizedCrop(size=(image_size, image_size), antialias=True),
7
+ v2.RandomHorizontalFlip(p=0.5),
8
+ v2.RandAugment(num_ops=2, magnitude=9),
9
+ v2.ToImage(),
10
+ v2.ToDtype(torch.float32, scale=True),
11
+ v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
12
+ ])
13
+
14
+ def get_eval_transforms(image_size=224):
15
+ return v2.Compose([
16
+ v2.Resize(size=(256, 256), antialias=True),
17
+ v2.CenterCrop(size=(image_size, image_size)),
18
+ v2.ToImage(),
19
+ v2.ToDtype(torch.float32, scale=True),
20
+ v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
21
+ ])
22
+
23
+ def get_batch_synthetics(num_classes):
24
+ return v2.RandomChoice([
25
+ v2.CutMix(num_classes=num_classes, alpha=1.0),
26
+ v2.MixUp(num_classes=num_classes, alpha=0.8)
27
+ ])
Version_2 - Copy/core/data_loader.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import pandas as pd
3
+ from torch.utils.data import DataLoader, Dataset
4
+ from PIL import Image
5
+ from core.augmentations import get_train_transforms, get_eval_transforms
6
+
7
+ class PhishingImageDataset(Dataset):
8
+ def __init__(self, dataframe, transform=None, image_col='image_path', label_col='label'):
9
+ self.dataframe = dataframe
10
+ self.transform = transform
11
+ self.image_col = image_col
12
+ self.label_col = label_col
13
+
14
+ def __len__(self):
15
+ return len(self.dataframe)
16
+
17
+ def __getitem__(self, idx):
18
+ row = self.dataframe.iloc[idx]
19
+ img_path = row[self.image_col]
20
+ label = row[self.label_col]
21
+
22
+ try:
23
+ image = Image.open(img_path).convert("RGB")
24
+ except Exception as e:
25
+ # Fallback to a blank image if file is missing/corrupted
26
+ print(f"Error loading {img_path}: {e}")
27
+ image = Image.new('RGB', (224, 224), (0, 0, 0))
28
+
29
+ if self.transform:
30
+ image = self.transform(image)
31
+
32
+ return image, label
33
+
34
+ def prepare_dataloaders(legit_csv_path, phishing_csv_path, batch_size=32, image_column_name='image_path'):
35
+ print(f"Sampling 5,000 rows from {legit_csv_path} and {phishing_csv_path}...")
36
+
37
+ df_legit = pd.read_csv(legit_csv_path).sample(n=5000, random_state=42)
38
+ df_legit['label'] = 0 # 0: Legit
39
+
40
+ df_phish = pd.read_csv(phishing_csv_path).sample(n=5000, random_state=42)
41
+ df_phish['label'] = 1 # 1: Phishing
42
+
43
+ df_all = pd.concat([df_legit, df_phish], ignore_index=True)
44
+ df_all = df_all.sample(frac=1, random_state=42).reset_index(drop=True)
45
+
46
+ # 80/10/10 Split -> 8000 Train, 1000 Val, 1000 Test
47
+ train_df = df_all.iloc[:8000]
48
+ val_df = df_all.iloc[8000:9000]
49
+ test_df = df_all.iloc[9000:]
50
+
51
+ train_ds = PhishingImageDataset(train_df, transform=get_train_transforms(), image_col=image_column_name)
52
+ val_ds = PhishingImageDataset(val_df, transform=get_eval_transforms(), image_col=image_column_name)
53
+ test_ds = PhishingImageDataset(test_df, transform=get_eval_transforms(), image_col=image_column_name)
54
+
55
+ train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=4)
56
+ val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=4)
57
+ test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=4)
58
+
59
+ return train_loader, val_loader, test_loader, 2
Version_2 - Copy/core/sam.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ class SAM(torch.optim.Optimizer):
4
+ def __init__(self, params, base_optimizer, rho=0.05, adaptive=False, **kwargs):
5
+ assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
6
+ defaults = dict(rho=rho, adaptive=adaptive, **kwargs)
7
+ super(SAM, self).__init__(params, defaults)
8
+ self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
9
+ self.param_groups = self.base_optimizer.param_groups
10
+ self.defaults.update(self.base_optimizer.defaults)
11
+
12
+ @torch.no_grad()
13
+ def first_step(self, zero_grad=False):
14
+ grad_norm = self._grad_norm()
15
+ for group in self.param_groups:
16
+ scale = group["rho"] / (grad_norm + 1e-12)
17
+ for p in group["params"]:
18
+ if p.grad is None: continue
19
+ self.state[p]["old_p"] = p.data.clone()
20
+ e_w = (torch.pow(p, 2) if group["adaptive"] else 1.0) * p.grad * scale.to(p)
21
+ p.add_(e_w)
22
+ if zero_grad: self.zero_grad()
23
+
24
+ @torch.no_grad()
25
+ def second_step(self, zero_grad=False):
26
+ for group in self.param_groups:
27
+ for p in group["params"]:
28
+ if p.grad is None: continue
29
+ p.data = self.state[p]["old_p"]
30
+ self.base_optimizer.step()
31
+ if zero_grad: self.zero_grad()
32
+
33
+ def _grad_norm(self):
34
+ shared_device = self.param_groups[0]["params"][0].device
35
+ norm = torch.norm(
36
+ torch.stack([
37
+ ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
38
+ for group in self.param_groups for p in group["params"]
39
+ if p.grad is not None
40
+ ]),
41
+ p=2
42
+ )
43
+ return norm
Version_2 - Copy/inference.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import timm
3
+ import ttach as tta
4
+ from PIL import Image
5
+ import os
6
+
7
+ from core.augmentations import get_eval_transforms
8
+
9
+ class ProductionAnalyzer:
10
+ def __init__(self, model_path, num_classes=2):
11
+ self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
+ self.transform = get_eval_transforms()
13
+ self.class_names = {0: "Legit", 1: "Phishing"}
14
+
15
+ print("Loading Production Model...")
16
+ base_model = timm.create_model('convnext_tiny', pretrained=False, num_classes=num_classes)
17
+
18
+ # Load weights and strip the 'module.' prefix caused by AveragedModel saving
19
+ state_dict = torch.load(model_path, map_location=self.device, weights_only=True)
20
+ clean_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
21
+
22
+ base_model.load_state_dict(clean_state_dict)
23
+ base_model.to(self.device)
24
+ base_model.eval()
25
+
26
+ # Wrap in TTA (Test Time Augmentation)
27
+ self.model = tta.ClassificationTTAWrapper(
28
+ base_model,
29
+ tta.aliases.five_crop_transform(224, 224)
30
+ )
31
+
32
+ def analyze_user_input(self, image_path):
33
+ if not os.path.exists(image_path):
34
+ return "Error: Image file not found."
35
+
36
+ image = Image.open(image_path).convert("RGB")
37
+ input_tensor = self.transform(image).unsqueeze(0).to(self.device)
38
+
39
+ with torch.no_grad():
40
+ logits = self.model(input_tensor)
41
+ probabilities = torch.softmax(logits, dim=1)
42
+ confidence, predicted_class = torch.max(probabilities, dim=1)
43
+
44
+ class_id = predicted_class.item()
45
+ return {
46
+ "prediction": self.class_names[class_id],
47
+ "class_id": class_id,
48
+ "confidence_score": f"{confidence.item() * 100:.2f}%"
49
+ }
50
+
51
+ if __name__ == "__main__":
52
+ # Ensure production_convnext_ema.pth exists in the directory before running
53
+ analyzer = ProductionAnalyzer(model_path="production_convnext_ema.pth")
54
+
55
+ # Example usage:
56
+ # result = analyzer.analyze_user_input("path_to_screenshot_to_test.jpg")
57
+ # print(result)
Version_2 - Copy/main.py ADDED
@@ -0,0 +1,175 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import sys
3
+ import time
4
+ import asyncio
5
+ import torch
6
+ import torch.nn.functional as F
7
+ import timm
8
+ from torchvision import transforms
9
+ from PIL import Image
10
+ from playwright.async_api import async_playwright
11
+
12
+ class ProductionAnalyzer:
13
+ def __init__(self, model_path="models/production_resnet_ema.pth", model_name="resnet18", num_classes=2):
14
+ """
15
+ Initializes the ResNet18 production model on available hardware.
16
+ """
17
+ self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
18
+ print(f"🖥️ Initializing backend hardware: {self.device}")
19
+
20
+ # 1. Initialize the ResNet18 architecture base
21
+ print(f"🏗️ Building network architecture: {model_name}...")
22
+ self.model = timm.create_model(model_name, pretrained=False, num_classes=num_classes)
23
+
24
+ # 2. Safety check for the weights file
25
+ if not os.path.exists(model_path):
26
+ raise FileNotFoundError(f"❌ Could not find weight file at: {model_path}\n"
27
+ f"Please ensure it is placed inside the 'models' folder.")
28
+
29
+ print(f"📥 Loading ResNet18 weights from {model_path}...")
30
+ state_dict = torch.load(model_path, map_location=self.device, weights_only=True)
31
+
32
+ # 3. 🛠️ FIXED: Strip wrapper prefixes and safely drop training metadata keys
33
+ clean_state_dict = {}
34
+ for key, value in state_dict.items():
35
+ if key == "n_averaged":
36
+ continue # Skip the training counter metadata so PyTorch doesn't throw an error
37
+
38
+ clean_key = key.replace('module.', '')
39
+ clean_state_dict[clean_key] = value
40
+
41
+ self.model.load_state_dict(clean_state_dict)
42
+
43
+ # 4. Lock the model for evaluation mode
44
+ self.model.eval()
45
+ self.model.to(self.device)
46
+ print("✅ ResNet18 Engine successfully loaded and locked for inference.")
47
+
48
+ # 5. Define standard normalizations expected by ResNet18
49
+ self.transform = transforms.Compose([
50
+ transforms.Resize((224, 224)),
51
+ transforms.ToTensor(),
52
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
53
+ ])
54
+
55
+ # Alphabetical class mapping (0: Legit, 1: Phishing)
56
+ self.class_names = ["Legit", "Phishing"]
57
+
58
+ def analyze_image(self, image_path):
59
+ """
60
+ Feeds the captured screenshot into the ResNet18 neural network.
61
+ """
62
+ try:
63
+ image = Image.open(image_path).convert('RGB')
64
+ input_tensor = self.transform(image).unsqueeze(0).to(self.device)
65
+
66
+ with torch.no_grad():
67
+ logits = self.model(input_tensor)
68
+ probabilities = F.softmax(logits[0], dim=0)
69
+ confidence, predicted_idx = torch.max(probabilities, 0)
70
+
71
+ return {
72
+ "prediction": self.class_names[predicted_idx.item()],
73
+ "confidence": f"{confidence.item() * 100:.2f}%"
74
+ }
75
+ except Exception as e:
76
+ return {"error": f"Model inference failed: {str(e)}"}
77
+
78
+
79
+ async def capture_screenshot(url, output_path="temp_inference.png"):
80
+ """
81
+ Launches a headless browser to safely capture a screenshot of the live URL.
82
+ """
83
+ if not url.startswith('http://') and not url.startswith('https://'):
84
+ url = 'https://' + url
85
+
86
+ print(f"🌐 Navigating to: {url} ...")
87
+
88
+ async with async_playwright() as p:
89
+ browser = await p.chromium.launch(headless=True)
90
+ context = await browser.new_context(
91
+ viewport={'width': 1280, 'height': 720},
92
+ ignore_https_errors=True, # Bypasses broken SSL certifications on malicious sites
93
+ accept_downloads=False
94
+ )
95
+
96
+ page = await context.new_page()
97
+ try:
98
+ # 12-second timeout to handle slow/malicious servers
99
+ await page.goto(url, timeout=12000, wait_until='domcontentloaded')
100
+ await asyncio.sleep(1) # Brief pause to let visual components load completely
101
+ await page.screenshot(path=output_path)
102
+ return output_path
103
+ except Exception as e:
104
+ print(f"❌ Failed to reach or capture the website: {e}")
105
+ return None
106
+ finally:
107
+ await page.close()
108
+ await browser.close()
109
+
110
+
111
+ async def main():
112
+ print("=" * 45)
113
+ print("🛡️ CHIMERA 2.0 LIVE URL DETECTOR ENGINE")
114
+ print("=" * 45)
115
+
116
+ try:
117
+ analyzer = ProductionAnalyzer()
118
+ except Exception as e:
119
+ print(e)
120
+ return
121
+
122
+ TEMP_IMG = "temp_inference.png"
123
+
124
+ try:
125
+ while True:
126
+ print("\n" + "-" * 45)
127
+ url_input = input("🔗 Enter URL to inspect (or type 'exit' to quit): ").strip()
128
+
129
+ if url_input.lower() == 'exit':
130
+ print("Shutting down engine...")
131
+ break
132
+ if not url_input:
133
+ continue
134
+
135
+ start_time = time.time()
136
+
137
+ # Step 1: Capture screenshot via Playwright
138
+ screenshot_file = await capture_screenshot(url_input, TEMP_IMG)
139
+
140
+ if screenshot_file and os.path.exists(screenshot_file):
141
+ # Step 2: Pass screenshot into ResNet18
142
+ print("🔍 Running Deep Learning Visual Inspection...")
143
+ result = analyzer.analyze_image(screenshot_file)
144
+
145
+ total_time = time.time() - start_time
146
+
147
+ # Step 3: Output results safely
148
+ print("\n" + "=" * 35)
149
+ print("📊 LIVE DETECTION REPORT")
150
+ print("=" * 35)
151
+ if "error" in result:
152
+ print(f"Result: {result['error']}")
153
+ else:
154
+ status_prefix = "🚨 ALERT!!" if result['prediction'] == "Phishing" else "✅ CLEAR:"
155
+ print(f"Verdict : {status_prefix} {result['prediction']}")
156
+ print(f"Confidence : {result['confidence']}")
157
+ print(f"Total Time : {total_time:.2f} seconds")
158
+ print("=" * 35)
159
+
160
+ if os.path.exists(TEMP_IMG):
161
+ os.remove(TEMP_IMG)
162
+ else:
163
+ print("❌ Inspection aborted. Visual fingerprint could not be gathered.")
164
+
165
+ except KeyboardInterrupt:
166
+ print("\nExiting execution gracefully...")
167
+ finally:
168
+ if os.path.exists(TEMP_IMG):
169
+ os.remove(TEMP_IMG)
170
+
171
+ if __name__ == "__main__":
172
+ if sys.platform == 'win32':
173
+ asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
174
+
175
+ asyncio.run(main())
Version_2 - Copy/models/production_resnet_ema.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:75e79d34179ef20f3fd3914e18c6867f4dfd356f8edd2b483eb5ba74779c7796
3
+ size 44794007