app.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import gradio as gr
from PIL import Image

# Model definition
class CyberNet(nn.Module):
    def __init__(self, input_channels=3, num_classes=87):
        super(CyberNet, self).__init__()
        self.conv1 = nn.Conv2d(input_channels, 93, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(93, 52, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(52, 172, kernel_size=5, stride=1, padding=1)
        self.pool = nn.MaxPool2d(2, 2)

        self.feature_size = self._get_conv_output_size(input_channels, (224, 224))
        self.fc1 = nn.Linear(self.feature_size, 512)
        self.fc2 = nn.Linear(512, num_classes)
        self.dropout = nn.Dropout(0.5)

    def _get_conv_output_size(self, input_channels, input_size):
        dummy_input = torch.zeros(1, input_channels, *input_size)
        x = self.pool(F.relu(self.conv1(dummy_input)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        return x.numel()

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# Load the model and weights
device = torch.device('cpu')
model = CyberNet(num_classes=87)
model.load_state_dict(torch.load('cyberNet.pt', map_location=device, weights_only=True))

# Data transformation pipeline
IMAGE_SIZE = 224  # Define the image size for resizing

data_transform = transforms.Compose([
    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),  # Resizing the input image to 224x224
    transforms.ToTensor(),  # Convert image to Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize based on ImageNet mean and std
])

# Define class names (you should replace this with your actual class names)
class_names = ['7ev3n', 'APosT', 'Adposhel', 'Agent', 'Agentb', 'Allaple', 'Alueron.gen!J', 'Amonetize', 'Androm', 'Bashlite', 'Bingoml', 'Blacksoul', 'BrowseFox', 'C2LOP.gen!g', 'Convagent', 'Copak', 'Delf', 'Dialplatform.B', 'Dinwod', 'Elex', 'Emotet', 'Escelar', 'Expiro', 'Fakerean', 'Fareit', 'Fasong', 'GandCrab', 'GlobelImposter', 'GootLoader', 'HLLP', 'HackKMS', 'Hlux', 'IcedId', 'Infy', 'Inject', 'Injector', 'InstallCore', 'KRBanker', 'Koadic', 'Kryptik', 'Kwampirs', 'Lamer', 'LemonDuck', 'Loki', 'Lolyda.AA1', 'Lolyda.AA2', 'Mimail', 'MultiPlug', 'Mydoom', 'Neoreklami', 'Neshta', 'NetWireRAT', 'Ngrbot', 'OnlinerSpambot', 'Orcus', 'Padodor', 'Plite', 'PolyRansom', 'QakBot', 'QtBot', 'Qukart', 'REvil', 'Ramdo', 'Regrun', 'Rekt Loader', 'Sakula', 'Salgorea', 'Scar', 'SelfDel', 'Small', 'Snarasite', 'Stantinko', 'Trickpak', 'Upantix', 'Upatre', 'VB', 'VBA', 'VBKrypt', 'VBNA', 'Vilsel', 'Vobfus', 'WBNA', 'Wecod', 'XTunnel', 'Zenpak', 'Zeus', 'benign']
print(class_names)

# Define the predict function
def predict(inp):
    inp = data_transform(inp).unsqueeze(0)  # Apply the transformation and add a batch dimension
    model.eval()
    with torch.no_grad():
        prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
        confidences = {class_names[i]: float(prediction[i]) for i in range(len(class_names))}
    return confidences

# Gradio app definition
gr.Interface(fn=predict,
             inputs=gr.Image(type="pil"),  # Accept PIL images
             outputs=gr.Label(num_top_classes=3),  # Output top 3 classes
             examples=["benign.png", "Upatre (88).png","Fakerean_aug_271_7932869.png" ,"Email-Worm.Win32.Mimail-q-9ee5ea94838d5af80a11d49a3c2e344c0122d227daa8e6f30e19858a91b3ffa1_RGB.png"]).launch(share=True)