Image-Based Malware Detection System

A web-based malware detection framework that converts executable files into grayscale image representations and uses Deep Convolutional Neural Networks (DCNNs) for multi-class malware classification.

Features

Binary-to-Image Conversion: Converts executables to grayscale images for visual pattern analysis
Deep Learning Classification: MobileNetV2-based model for multi-class malware family detection
High Accuracy: 92% accuracy with 31 malware families (Blended Malware Image Dataset)
Grad-CAM Explainability: Visual explanations showing which binary regions influenced the decision
Real-time Analysis: FastAPI-powered backend for quick file scanning
Forensic Reports: Generate PDF and JSON reports for analysis documentation
Scan History: Track and review all previous scans with statistics
Modern UI: Dark-themed responsive interface built with TailwindCSS

Supported Datasets

Dataset	Classes	Samples	Expected Accuracy	Download Size	Recommended
Blended Malware Image	31	27,494	92%	~2 GB	⭐ Current

Tech Stack

Backend: Python 3.9+, FastAPI, SQLAlchemy, SQLite
ML: PyTorch, MobileNetV2, Grad-CAM
Frontend: HTML, TailwindCSS, Vanilla JavaScript
Training: Local GPU or Google Colab
Package Manager: uv (recommended)

Project Structure

malware-detector/
├── backend/
│   ├── app/
│   │   ├── main.py              # FastAPI application entry point
│   │   ├── config.py            # Configuration settings
│   │   ├── database.py          # SQLAlchemy database setup
│   │   ├── models/
│   │   │   ├── db_models.py     # Database models
│   │   │   └── schemas.py       # Pydantic schemas
│   │   ├── routers/
│   │   │   ├── scan.py          # File scanning endpoints
│   │   │   ├── history.py       # Scan history endpoints
│   │   │   └── reports.py       # Report generation endpoints
│   │   ├── services/
│   │   │   ├── classifier.py    # ML model inference
│   │   │   ├── binary_to_image.py  # Binary to image conversion
│   │   │   ├── gradcam.py       # Grad-CAM visualization
│   │   │   └── report_generator.py # PDF/JSON report generation
│   │   └── utils/
│   │       ├── security.py      # File validation
│   │       └── helpers.py       # Utility functions
│   ├── ml_model/
│   │   ├── model.pth            # Trained PyTorch model weights
│   │   ├── class_labels.json    # Active class label mappings
│   │   ├── class_labels_malimg.json   # Malimg labels (25 classes)
│   │   └── class_labels_big2015.json  # BIG 2015 labels (10 classes)
│   ├── static/
│   │   ├── gradcam/             # Generated visualizations
│   │   └── reports/             # Generated reports
│   ├── uploads/                 # Temporary file storage
│   ├── requirements.txt         # Python dependencies
│   └── malware_detector.db      # SQLite database
├── frontend/
│   ├── index.html               # Main web interface
│   ├── css/styles.css           # Stylesheet
│   └── js/app.js                # Frontend JavaScript
├── training/
│   ├── train_model.ipynb        # Google Colab training notebook
│   ├── prepare_malimg.py        # Malimg dataset preparation (recommended)
│   ├── prepare_big2015.py       # Microsoft BIG 2015 dataset preparation
│   └── data_preparation.py      # General dataset preparation script
├── datasets/                    # Training datasets (not included)
└── tests/

Prerequisites

Python 3.11+ (required)
uv - Python package manager (install guide)
Web browser (Chrome, Firefox, Edge)

Quick Start

1. Install Dependencies

# Navigate to the project directory
cd malware-detector

# Install all dependencies using uv
uv sync

This creates a virtual environment in .venv/ and installs all required packages.

2. Configure Environment

# Copy the example environment file
cp backend/.env.example backend/.env

Edit backend/.env if needed (defaults work for local development).

3. Verify Model Files

Ensure these files exist in backend/ml_model/:

backend/ml_model/
├── model.pth            # Trained model weights (~11 MB)
├── class_labels.json    # Class labels (31 malware families)
└── training_info.json   # Training metadata

If not present, train the model using the training notebook (see Training section).

Running the Application

Option 1: Quick Start (Two Terminals)

Terminal 1 - Start Backend:

cd malware-detector
uv run uvicorn backend.app.main:app --reload --port 8000

Terminal 2 - Start Frontend:

cd frontend
python -m http.server 3000

Access the application:

Frontend UI: http://localhost:3000
API Documentation: http://localhost:8000/docs
API ReDoc: http://localhost:8000/redoc

Option 2: Using VS Code

Open the project in VS Code
Open two terminals
Run the backend and frontend commands as shown above

Option 3: Production Mode

uv run uvicorn backend.app.main:app --host 0.0.0.0 --port 8000 --workers 4

Usage

Web Interface

Open http://127.0.0.1:5500 in your browser
Click "Select File" or drag-and-drop a file
Supported formats: .exe, .dll, .sys, .elf, .so, .bin
Click "Scan File" to analyze
View results including:
- Classification (Malware/Benign)
- Confidence score
- Malware family name
- Grayscale visualization
- Grad-CAM heatmap

API Usage

Scan a file:

curl -X POST "http://127.0.0.1:8000/api/scan/" -F "file=@sample.exe"

Get scan history:

curl "http://127.0.0.1:8000/api/history/"

Get statistics:

curl "http://127.0.0.1:8000/api/history/stats"

Download PDF report:

curl "http://127.0.0.1:8000/api/reports/1/pdf" -o report.pdf

API Endpoints

Endpoint	Method	Description
`/api/scan/`	POST	Upload and scan a file
`/api/scan/{id}`	GET	Get scan result by ID
`/api/scan/{id}`	DELETE	Delete a scan record
`/api/history/`	GET	List all scans (paginated)
`/api/history/stats`	GET	Get scanning statistics
`/api/history/search`	GET	Search scans by filename/hash
`/api/reports/{id}/pdf`	GET	Download PDF report
`/api/reports/{id}/json`	GET	Get JSON report

Supported File Types

Extension	Type	Description
`.exe`	PE	Windows Executable
`.dll`	PE	Windows Dynamic Link Library
`.sys`	PE	Windows System Driver
`.elf`	ELF	Linux Executable
`.so`	ELF	Linux Shared Object
`.bin`	Binary	Generic Binary File

Maximum file size: 50 MB

Model Information

Architecture

Base Model: MobileNetV2 (pretrained on ImageNet, fine-tuned)
Classifier Head: Custom multi-layer classifier
- Dropout(0.3) → Linear(1280, 512) → ReLU → BatchNorm → Dropout(0.2) → Linear(512, num_classes)
Input Size: 256x256 RGB (grayscale converted to 3-channel)
Output: 31 malware family classes
Model Size: ~11 MB

Current Model: Blended Malware Image Dataset (31 families)

Class ID	Family Name	Precision	Recall
0	Adposhel	98%	100%
1	Agent	84%	78%
2	Allaple	94%	89%
3	Alueron.gen!J	100%	100%
4	Amonetize	98%	88%
5	Androm	88%	99%
6	Autorun	78%	89%
7	BrowseFox	94%	94%
8	C2LOP.gen!g	100%	100%
9	Dialplatform.B	100%	100%
10	Dinwod	95%	97%
11	Elex	93%	95%
12	Expiro	81%	77%
13	Fakerean	100%	100%
14	Fasong	100%	100%
15	HackKMS	97%	98%
16	Hlux	100%	99%
17	Injector	95%	78%
18	InstallCore	99%	97%
19	Lolyda.AA1	100%	100%
20	Lolyda.AA2	100%	100%
21	MultiPlug	98%	90%
22	Neoreklami	96%	97%
23	Neshta	73%	61%
24	Regrun	100%	100%
25	Sality	57%	76%
26	Snarasite	99%	100%
27	Stantinko	96%	94%
28	VBA	99%	100%
29	VBKrypt	90%	93%
30	Vilsel	99%	100%

Overall Accuracy: 92.19%

Class Labels Format

The class_labels.json file maps class indices to names:

{
  "0": {"name": "Adposhel", "is_malware": true},
  "1": {"name": "Agent", "is_malware": true},
  "2": {"name": "Allaple", "is_malware": true},
  ...
}

Note: This dataset contains only malware families (no benign class).

Training Your Own Model

Recommended: Blended Malware Image Dataset (92% accuracy, Current Model)

Step 1: Download Dataset

Create a Kaggle account at https://www.kaggle.com
Go to: https://www.kaggle.com/datasets/gauravpendharkar/blended-malware-image-dataset
Download blended-malware-image-dataset.zip (~2 GB)

Or via CLI:

kaggle datasets download -d gauravpendharkar/blended-malware-image-dataset

Step 2: Prepare Dataset

# Extract the archive
unzip blended-malware-image-dataset.zip -d malware_dataset

The dataset contains 31 malware families with pre-split train/val:

malware_dataset/
├── train/                 (19,736 images, 72%)
│   ├── Adposhel/         (700 samples)
│   ├── Agent/            (700 samples)
│   ├── Allaple/          (700 samples)
│   ├── ... (28 more families)
│   └── Vilsel/           (700 samples)
└── val/                   (7,758 images, 28%)
    ├── Adposhel/         (288 samples)
    ├── ...
    └── Vilsel/           (292 samples)
Total: 27,494 samples across 31 classes

Step 3: Train Locally (GPU Required)

cd training
pip install -r requirements.txt

# Install PyTorch with CUDA
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

Open train_malware_detector.ipynb in Jupyter and run all cells.

Step 4: Deploy Model

The trained model is automatically saved to backend/ml_model/:

model.pth - Model weights (~11 MB)
class_labels.json - Class mappings
training_info.json - Training metadata

Restart the backend server to load the new model.

Alternative: Malimg Dataset (94-98% accuracy)

Download from: https://www.kaggle.com/datasets/keerthicheepurupalli/malimg-dataset
Contains 25 malware families, 9,339 samples
Smaller and faster to train

Alternative: Microsoft BIG 2015 Dataset (97-99% accuracy)

For highest accuracy but larger download (~37 GB compressed, ~500 GB extracted):

# Download from Kaggle
kaggle competitions download -c malware-classification

# Extract and organize
7z x train.7z
cd training
python prepare_big2015.py --train-dir ../train --labels ../trainLabels.csv --output ../microsoft-big2015

Custom Dataset

Organize your dataset with folders per class:

dataset/
├── MalwareFamily1/
│   ├── sample1.png  # or .bytes, .exe
│   └── sample2.png
├── MalwareFamily2/
│   └── ...
└── Benign/
    └── ...

Training Dependencies (Optional)

If you want to train or fine-tune the model, install the training extras:

uv sync --extra training

This adds matplotlib, seaborn, scikit-learn, jupyter, ipywidgets, and tqdm.

Running Training Notebook

uv run jupyter notebook training/train_malware_detector.ipynb

Note for GPU Training: The default PyTorch installation is CPU-only. For CUDA support:

# CUDA 12.1
uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

# CUDA 11.8
uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

Development

Running Tests

uv run pytest

Code Formatting & Linting

# Check code style
uv run ruff check .

# Auto-fix issues
uv run ruff check . --fix

Troubleshooting

Common Issues

1. "Module not found" or import errors:

# Make sure you're using uv run
uv run python -c "import torch; print(torch.__version__)"

2. Port already in use:

# Change the port number
uv run uvicorn backend.app.main:app --reload --port 8001

3. Model not loading:

Verify model.pth exists in backend/ml_model/
Check file size (should be ~12 MB)
Ensure class_labels.json has correct format

4. CORS errors:

Ensure frontend is served via HTTP server (not file://)
Backend CORS is configured for all origins in development (allow_origins=["*"])

5. Low confidence predictions:

Model shows actual predicted class regardless of confidence
Low confidence (<45%) marks file as safe
May indicate file type not in training data

Performance

Inference Time: ~100-500ms per file (CPU), ~50-100ms (GPU)
Model Size: ~12 MB
Memory Usage: ~500 MB RAM

Security Considerations

Files are validated before processing
Uploaded files are automatically deleted after scanning
Maximum file size limits enforced
Input sanitization prevents path traversal attacks
SQLite database with parameterized queries

Warning: Handle malware samples in isolated environments only (VM/sandbox).

Project Status

License

For educational and research purposes only.

Acknowledgments

MobileNetV2: https://arxiv.org/abs/1801.04381
Grad-CAM: https://arxiv.org/abs/1610.02391
Malware Visualization: https://arxiv.org/abs/1703.00138

Name		Name	Last commit message	Last commit date
Latest commit History 9 Commits
backend		backend
frontend		frontend
training		training
.gitignore		.gitignore
README.md		README.md
malware_detector.db		malware_detector.db
output.png		output.png
pyproject.toml		pyproject.toml
setup.md		setup.md
uv.lock		uv.lock

Folders and files

Latest commit

History

Repository files navigation