🦠 MalVision

Deep Learning Malware Detection via Binary Visualization

Transform malware binaries into images, then let neural networks see what humans can't.

Getting Started • Results • Architecture • Dataset

---

What is This?

MalVision converts malware executables into grayscale images and uses Convolutional Neural Networks to classify them into malware families. This technique leverages the visual patterns that emerge when binary data is rendered as pixels – patterns that are surprisingly distinctive for different malware variants.

┌─────────────────┐     ┌─────────────────┐     ┌─────────────────┐
│  Malware Binary │ ──► │ Grayscale Image │ ──► │  CNN Classifier │
│   (bytes)       │     │  (visualization)│     │  (25 families)  │
└─────────────────┘     └─────────────────┘     └─────────────────┘

Features

• Transfer Learning - ResNet18 & EfficientNet-B0 pretrained on ImageNet
• Custom CNN - Simple architecture for educational purposes
• Data Augmentation - Rotation, flipping, color jitter for robust training
• Rich Visualization - Training curves, confusion matrices, per-class analysis
• GPU Accelerated - CUDA support for fast training
• Model Checkpointing - Save best models with full training state

Quick Start

Installation

We use Conda for environment management and uv for fast package installation.

# Clone the repository
git clone https://github.com/Y4NN777/MalVision.git
cd MalVision

# 1. Create Conda environment
conda create -n malvision python=3.10 -y
conda activate malvision

# 2. Install uv (fast installer)
pip install uv

# 3. Install dependencies from pyproject.toml
uv pip install -r pyproject.toml

Download Dataset

Download the MalImg Dataset from Kaggle:

# Using helper script
python scripts/fetch_dataset.py

# OR manually download/unzip 'malimg-dataset.zip' to 'malware_data/' folder

Run the Notebook

# Start Jupyter
jupyter notebook malware_detection.ipynb

The notebook guides you through:

1. Dataset exploration
2. Data preparation with augmentation
3. Model training (ResNet18)
4. Evaluation with confusion matrix
5. Single image prediction
6. Model export

Project Structure

MalVision/
├── malware_detection.ipynb   # Main notebook (run this!)
├── pyproject.toml            # Dependencies & Config
├── utils/
│   ├── __init__.py           # Package exports
│   └── helpers.py            # Dataset, models, visualization
├── scripts/
│   └── fetch_dataset.py      # Dataset helper
├── malware_data/             # Dataset folder (create this!)
├── .gitignore
├── LICENSE
└── README.md

Architecture

Model Options

Model	Parameters	Best For
SimpleCNN	~2M	Learning & experimentation
ResNet18	~11M	Balance of speed & accuracy
EfficientNet-B0	~5M	Best accuracy

SimpleCNN Architecture

Input (3, 224, 224)
    │
    ├── Conv2d(3→32) + BN + ReLU + MaxPool  →  (32, 112, 112)
    ├── Conv2d(32→64) + BN + ReLU + MaxPool →  (64, 56, 56)
    ├── Conv2d(64→128) + BN + ReLU + MaxPool → (128, 28, 28)
    ├── Conv2d(128→256) + BN + ReLU + MaxPool → (256, 14, 14)
    │
    ├── Flatten → (256 × 14 × 14 = 50,176)
    ├── Dropout(0.5) + Linear(50176→512) + ReLU
    └── Dropout(0.5) + Linear(512→25)
    
Output: 25 classes (malware families)

Dataset

MalImg Dataset - 9,339 malware samples across 25 families

Family	Samples		Family	Samples
Adialer.C	122		Lolyda.AA1	213
Agent.FYI	116		Lolyda.AA2	184
Allaple.A	2,949		Lolyda.AA3	123
Allaple.L	1,591		Lolyda.AT	159
Alueron.gen!J	198		Malex.gen!J	136
Autorun.K	106		Obfuscator.AD	142
C2LOP.gen!g	200		Rbot!gen	158
C2LOP.P	146		Skintrim.N	80
Dialplatform.B	177		Swizzor.gen!E	128
Dontovo.A	162		Swizzor.gen!I	132
Fakerean	381		VB.AT	408
Instantaccess	431		Wintrim.BX	97
Yuner.A	800

Results

Expected performance with default configuration:

Metric	Value
Test Accuracy	~97%
Training Time	~15 min (GPU)
Best Epoch	~12-15

Results & Models

Training artifacts are saved locally:

• training_curves.png: Visual training history
• confusion_matrix.png: Per-class performance
• classification_report.txt: F1-scores and precision/recall

📂 Download Trained Models

Models are too large for Git. Download checkpoints from: ☁️ Google Drive Link

Place downloaded .pth files in the models/ directory:

• best_model.pth: Highest validation accuracy
• malware_model_complete.pth: Full checkpoint with training history

Training Configuration

Parameter	Default	Description
epochs	20	Training epochs
batch_size	32	Samples per batch
learning_rate	0.001	Initial LR
weight_decay	1e-4	L2 regularization
train_split	0.70	Training data ratio
val_split	0.15	Validation data ratio
patience	5	Early stopping patience

Tech Stack

• PyTorch - Deep learning framework
• torchvision - Pretrained models & transforms
• NumPy & Pandas - Data manipulation
• Matplotlib & Seaborn - Visualization
• scikit-learn - Metrics & analysis
• PIL - Image processing
• tqdm - Progress bars

Educational Context

This project was developed as part of a Deep Learning course, demonstrating:

1. Image Classification - End-to-end CNN pipeline
2. Transfer Learning - Leveraging pretrained models
3. Data Augmentation - Preventing overfitting
4. Model Evaluation - Comprehensive metrics analysis
5. Visualization - Interpreting model behavior

Acknowledgments

• MalImg Dataset - Malware visualization dataset
• PyTorch - Deep learning framework
• Original research: "Malware Images: Visualization and Automatic Classification" by Nataraj et al.

---