Vision-Transformer-ViT-for-Image-Classification

KO | EN

gitlite — search

#javascript #python #hacktoberfest #react #ai #typescript #llm #go #golang #android #machine-learning #rust #deep-learning #linux

Vision-Transformer-ViT-for-Image-Classification

★ 13 Open GitHub ↗

No description available.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Vision Transformer (ViT) for Cats vs Dogs Image Classification

This repository contains a **Vision Transformer (ViT)** based image classifier trained to distinguish between **Cats vs Dogs**. The implementation uses PyTorch and a lightweight **pre-trained ViT** backbone (`vit_tiny_patch16_224`) from the `timm` library and includes a simple **Streamlit demo** for deployment.

---

## 📌 Project Overview

- **Task**: Binary image classification (Cat vs Dog)
- **Model**: Vision Transformer (**ViT-Tiny**, patch size 16, input 224×224)
- **Framework**: PyTorch
- **Pretraining**: ImageNet (via `timm`)
- **Dataset**: Microsoft/Kaggle Cats vs Dogs (folder: `PetImages/`)
- **Best Test Accuracy**: **94.53%**

---

## 📊 Dataset

This project uses the well-known **Cats vs Dogs** dataset (commonly distributed with a `PetImages/` directory containing `Cat/` and `Dog/` subfolders).

✅ **Dataset sources**:
- Microsoft official download: https://www.microsoft.com/en-us/download/details.aspx?id=54765 ([microsoft.com](https://www.microsoft.com/en-us/download/details.aspx?id=54765&utm_source=chatgpt.com))
- Kaggle competition page: https://www.kaggle.com/c/dogs-vs-cats ([kaggle.com](https://www.kaggle.com/c/dogs-vs-cats?utm_source=chatgpt.com))

### Expected Folder Structure

Place the dataset in the project root like this:

```
./PetImages/
   ├── Cat/
   └── Dog/
```

---

## 🧹 Preprocessing

Images are preprocessed using standard ImageNet normalization:

- Resize to **224 × 224**
- Convert to tensor
- Normalize using:
  - mean = `[0.485, 0.456, 0.406]`
  - std  = `[0.229, 0.224, 0.225]`

---

## 🧠 Model

### Backbone
- `timm.create_model('vit_tiny_patch16_224', pretrained=True, num_classes=2)`

---

## 🔧 Training Configuration

| Setting | Value |
|---|---|
| Input Size | 224 × 224 |
| Train / Val / Test Split | 70% / 15% / 15% |
| Batch Size | 32 |
| Epochs | 10 |
| Optimizer | Adam |
| Learning Rate | 1e-4 |
| Loss Function | CrossEntropyLoss |
| Device | CUDA (if available) |

---

## 📈 Results

Final performance from the training run:

| Metric | Value |
|---|---|
| **Train Accuracy (Epoch 10)** | **98.60%** |
| **Validation Accuracy (Epoch 10)** | **94.43%** |
| **Test Accuracy** | **94.53%** |
| Test Loss | 0.1571 |

The model performs strongly due to ViT pretraining and effective fine-tuning on the Cats vs Dogs dataset.

---

## 📉 Visualizations

The notebook includes training curves:
- Training vs Validation **Loss**
- Training vs Validation **Accuracy**

These help analyze overfitting, convergence, and generalization.

---

## 💾 Checkpointing

The best model is saved as:

```
checkpoint.pt
```

This checkpoint is reused for inference and for the Streamlit demo.

---

## ▶️ How to Run

```bash
# Clone the repository
git clone https://github.com/<your-username>/vit-cats-vs-dogs.git
cd vit-cats-vs-dogs

# Install dependencies
pip install -r requirements.txt

# Run the notebook
jupyter notebook ViT_for_Image_Classification.ipynb
```

---

## 🚀 Deployment (Streamlit Demo)

A minimal Streamlit app is included to upload an image and get predictions.

Run:

```bash
streamlit run vit_bonus_demo.py
```

> The demo loads `checkpoint.pt` and performs inference on user-uploaded images.

---

## 📁 Project Structure

```
.
├── ViT_for_Image_Classification.ipynb
├── README.md
├── requirements.txt
├── checkpoint.pt
├── vit_bonus_demo.py

```

---

## 📚 References

1. Dosovitskiy et al., **An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale** (ICLR 2021)
2. `timm` library (PyTorch Image Models): https://github.com/huggingface/pytorch-image-models
3. Cats vs Dogs dataset (Microsoft): https://www.microsoft.com/en-us/download/details.aspx?id=54765 ([microsoft.com](https://www.microsoft.com/en-us/download/details.aspx?id=54765&utm_source=chatgpt.com))

---

## 🙌 Acknowledgements

This project was completed as part of **CSE 676-B: Deep Learning** at the University at Buffalo by Poojan Kaneriya (poojanra@buffalo.edu).

---

⭐ If this project helped you, feel free to star the repository!

Vision-Transformer-ViT-for-Image-Classification

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README