Spaces:
Runtime error
Runtime error
File size: 2,524 Bytes
b3a532e | 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 | import torch
import pandas as pd
from torch.utils.data import DataLoader, Dataset
from PIL import Image
from core.augmentations import get_train_transforms, get_eval_transforms
class PhishingImageDataset(Dataset):
def __init__(self, dataframe, transform=None, image_col='image_path', label_col='label'):
self.dataframe = dataframe
self.transform = transform
self.image_col = image_col
self.label_col = label_col
def __len__(self):
return len(self.dataframe)
def __getitem__(self, idx):
row = self.dataframe.iloc[idx]
img_path = row[self.image_col]
label = row[self.label_col]
try:
image = Image.open(img_path).convert("RGB")
except Exception as e:
# Fallback to a blank image if file is missing/corrupted
print(f"Error loading {img_path}: {e}")
image = Image.new('RGB', (224, 224), (0, 0, 0))
if self.transform:
image = self.transform(image)
return image, label
def prepare_dataloaders(legit_csv_path, phishing_csv_path, batch_size=32, image_column_name='image_path'):
print(f"Sampling 5,000 rows from {legit_csv_path} and {phishing_csv_path}...")
df_legit = pd.read_csv(legit_csv_path).sample(n=5000, random_state=42)
df_legit['label'] = 0 # 0: Legit
df_phish = pd.read_csv(phishing_csv_path).sample(n=5000, random_state=42)
df_phish['label'] = 1 # 1: Phishing
df_all = pd.concat([df_legit, df_phish], ignore_index=True)
df_all = df_all.sample(frac=1, random_state=42).reset_index(drop=True)
# 80/10/10 Split -> 8000 Train, 1000 Val, 1000 Test
train_df = df_all.iloc[:8000]
val_df = df_all.iloc[8000:9000]
test_df = df_all.iloc[9000:]
train_ds = PhishingImageDataset(train_df, transform=get_train_transforms(), image_col=image_column_name)
val_ds = PhishingImageDataset(val_df, transform=get_eval_transforms(), image_col=image_column_name)
test_ds = PhishingImageDataset(test_df, transform=get_eval_transforms(), image_col=image_column_name)
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=4)
test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=4)
return train_loader, val_loader, test_loader, 2 |