AventIQ-AI/Single-Label-General-Image-Classifier

🧠 Image Classification AI Model (CIFAR-100)

This repository contains a Vision Transformer (ViT)-based AI model fine-tuned for image classification on the CIFAR-100 dataset. The model is built using google/vit-base-patch16-224, quantized to FP16 for efficient inference, and delivers high accuracy in multi-class image classification tasks.

---

🚀 Features

• 🖼️ Task: Image Classification
• 🧠 Base Model: google/vit-base-patch16-224 (Vision Transformer)
• 🧪 Quantized: FP16 for faster and memory-efficient inference
• 🎯 Dataset: 100 fine-grained object categories
• ⚡ CUDA Enabled: Optimized for GPU acceleration
• 📈 High Accuracy: Fine-tuned and evaluated on validation split

---

📊 Dataset Used

Hugging Face Dataset: tanganke/cifar100

• Description: CIFAR-100 is a dataset of 60,000 32×32 color images in 100 classes (600 images per class)
• Split: 50,000 training images and 10,000 test images
• Categories: Animals, Vehicles, Food, Household items, etc.
• License: MIT License (from source)

from datasets import load_dataset

dataset = load_dataset("tanganke/cifar100")

🛠️ Model & Training Configuration

• Model: google/vit-base-patch16-224

• Image Size: 224x224 (resized from 32x32)

• Framework: Hugging Face Transformers & Datasets

• Training Environment: Kaggle Notebook with CUDA

• Epochs: 5–10 (with early stopping)

• Batch Size: 32

• Optimizer: AdamW

• Loss Function: CrossEntropyLoss

✅ Evaluation & Scoring

• Accuracy: ~70–80% (varies by configuration)

• Validation Tool: evaluate or sklearn.metrics

• Metric: Accuracy, Top-1 and Top-5 scores

• Inference Speed: Significantly faster after quantizationextractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")

🔍 Inference Example

from PIL import Image
import torch

def predict(image_path):
    image = Image.open(image_path).convert("RGB")
    inputs = feature_extractor(images=image, return_tensors="pt").to("cuda")
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = logits.argmax(-1).item()
    return dataset["train"].features["fine_label"].int2str(predicted_class)

print(predict("sample_image.jpg"))

📁 Folder Structure

📦image-classification-vit ┣ 📂vit-cifar100-fp16 ┣ 📜train.py ┣ 📜inference.py ┣ 📜README.md ┗ 📜requirements.txt