univeral-img-classifier-small

Universal Image Classifier 🖼️🤖

A flexible, PyTorch-based image classification framework designed for easy adaptation to any image classification task. Perfect for researchers, practitioners, and anyone looking to quickly prototype and deploy image classification models.

🚀 Features

• 🏗️ Multiple Architectures: Linear, MLP, Deep MLP, and Residual MLP models
• ⚙️ Flexible Configuration: Easy hyperparameter tuning via config files
• 🔮 Universal Inference: Works with any trained model and dataset
• 🚀 Production Ready: Optimized for both training and deployment
• 📊 Modern Training: Data augmentation, early stopping, learning rate scheduling
• 🤗 Hugging Face Ready: Easy integration with Hugging Face Hub
• 📈 Comprehensive Evaluation: Built-in metrics, visualization, and model analysis
• 📦 Easy Dataset Loading: Support for folder structure and CSV datasets

📊 Model Architectures

🔧 Installation

# Clone the repository
git clone https://github.com/hightechguys79/universal-image-classifier.git
cd universal-image-classifier

# Install dependencies
pip install -r requirements.txt

🎯 Quick Start

1. Training a Model

from train import UniversalTrainer
from config import ModelConfig, TrainingConfig, DataConfig
from utils import load_dataset_from_folder

# Load your dataset
image_paths, labels, class_names = load_dataset_from_folder("path/to/your/data")

# Configure model
model_config = ModelConfig(
    num_classes=len(class_names),
    hidden_dim=256,
    num_layers=6,
    dropout_rate=0.1
)

training_config = TrainingConfig(
    batch_size=64,
    learning_rate=0.001, 
    num_epochs=100
)

# Train the model
trainer = UniversalTrainer(model_config, training_config, DataConfig())
trainer.train(train_dataset, val_dataset, model_name="mlp_deep_residual")

2. Making Predictions

from inference import load_classifier

# Load your trained model
classifier = load_classifier(
    model_path="outputs/mlp_deep_residual_best.pth",
    config_path="outputs/mlp_deep_residual_config.json",
    class_names=["cat", "dog", "bird", "fish"]
)

# Predict single image
result = classifier.predict("test_image.jpg")
print(f"Predicted: {result['predicted_class']} ({result['confidence']:.3f})")

# Get top-3 predictions
top_3 = classifier.get_top_k_predictions("test_image.jpg", k=3)
for pred in top_3:
    print(f"{pred['class_name']}: {pred['probability']:.3f}")

3. Running Examples

# Train a demo model
python examples/train_example.py

# Test inference
python examples/inference_example.py

# Prepare for Hugging Face
python examples/huggingface_integration.py

📁 Project Structure

universal-image-classifier/
├── README.md                    # This file
├── requirements.txt             # Dependencies
├── config.py                   # Configuration classes
├── models.py                   # Model architectures
├── train.py                    # Training pipeline
├── inference.py                # Inference pipeline
├── utils.py                    # Utility functions
├── examples/                   # Example scripts
│   ├── train_example.py        # Training example
│   ├── inference_example.py    # Inference example
│   └── huggingface_integration.py # HF Hub integration
└── assets/                     # Documentation assets

🎨 Supported Dataset Formats

1. Folder Structure

your_dataset/
├── class1/
│   ├── img1.jpg
│   └── img2.jpg
└── class2/
    ├── img3.jpg
    └── img4.jpg

2. CSV Format

image_path,label
path/to/img1.jpg,cat
path/to/img2.jpg,dog

⚙️ Configuration

Model Configuration

ModelConfig(
    input_height=64,        # Input image height
    input_width=64,         # Input image width
    num_classes=10,         # Number of output classes
    hidden_dim=128,         # Hidden layer dimension
    num_layers=4,           # Number of layers
    dropout_rate=0.1,       # Dropout rate
    use_batch_norm=True     # Use batch normalization
)

Training Configuration

TrainingConfig(
    batch_size=32,                    # Batch size
    learning_rate=0.001,              # Learning rate
    num_epochs=100,                   # Maximum epochs
    weight_decay=1e-4,                # L2 regularization
    early_stopping_patience=10,       # Early stopping patience
    validation_split=0.2              # Validation split ratio
)

🔬 Advanced Usage

Custom Dataset Integration

from torch.utils.data import Dataset
from utils import load_dataset_from_csv

# Load from CSV
image_paths, labels, class_names = load_dataset_from_csv(
    "dataset.csv", 
    image_col="image_path", 
    label_col="label"
)

# Create custom dataset
class MyDataset(Dataset):
    def __init__(self, image_paths, labels, transform=None):
        self.image_paths = image_paths
        self.labels = labels
        self.transform = transform
    
    def __getitem__(self, idx):
        # Your custom loading logic
        pass

Model Evaluation

from utils import evaluate_model, plot_confusion_matrix

# Evaluate model
metrics = evaluate_model(model, test_loader, class_names, device)
print(f"Accuracy: {metrics['accuracy']:.4f}")

# Plot confusion matrix
plot_confusion_matrix(
    metrics['true_labels'], 
    metrics['predictions'], 
    class_names
)

Hyperparameter Optimization with Weights & Biases

import wandb

# Initialize sweep
sweep_config = {
    'method': 'random',
    'parameters': {
        'hidden_dim': {'values': [128, 256, 512]},
        'learning_rate': {'min': 0.0001, 'max': 0.01},
        'num_layers': {'values': [3, 4, 5, 6]},
        'dropout_rate': {'min': 0.0, 'max': 0.3}
    }
}

sweep_id = wandb.sweep(sweep_config, project="universal-classifier")
wandb.agent(sweep_id, train_function)

🤗 Hugging Face Integration

Export Model

from utils import export_to_huggingface

export_to_huggingface(
    model_path="outputs/model.pth",
    config_path="outputs/config.json",
    class_names=["cat", "dog"],
    model_name="mlp_deep_residual",
    output_dir="huggingface_model"
)

Upload to Hub

from huggingface_hub import upload_folder

upload_folder(
    folder_path="huggingface_model",
    repo_id="your-username/universal-classifier",
    repo_type="model"
)

📈 Performance Tips

1. 🎯 Data Augmentation: Enable for small datasets
2. 📊 Learning Rate: Start with 0.001, reduce if loss plateaus
3. 🏗️ Architecture: Use residual connections for deep networks
4. 🎛️ Regularization: Adjust dropout and weight decay based on overfitting
5. 📦 Batch Size: Larger batches for stable training, smaller for better generalization
6. ⏰ Early Stopping: Prevent overfitting with patience-based stopping

🔍 Model Analysis

from models import count_parameters, calculate_model_size_mb

model = create_model("mlp_deep_residual", config)
print(f"Parameters: {count_parameters(model):,}")
print(f"Model size: {calculate_model_size_mb(model):.2f} MB")

🎨 Example Applications

• 🏠 Object Recognition: Classify everyday objects
• 🏥 Medical Imaging: Diagnostic image classification
• 🏭 Quality Control: Defect detection in manufacturing
• 🛡️ Content Moderation: Automatic image filtering
• 🦎 Wildlife Monitoring: Species identification
• 🎨 Art Classification: Style and genre recognition
• 🍕 Food Recognition: Cuisine and dish classification

📊 Benchmarks

Benchmarks run on NVIDIA RTX 3080, batch size 64

🤝 Contributing

We welcome contributions! Please see our Contributing Guidelines for details.

Development Setup

# Clone and install in development mode
git clone https://github.com/hightechguys79/universal-image-classifier.git
cd universal-image-classifier
pip install -e .

# Run tests
python -m pytest tests/

# Format code
black .
isort .

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

🙏 Acknowledgments

• PyTorch team for the excellent deep learning framework
• Hugging Face for the model hub infrastructure
• The open-source community for inspiration and feedback

📞 Support

• 🐛 Bug Reports: GitHub Issues
• 💬 Discussions: GitHub Discussions
• 📧 Email: [email protected]

🗺️ Roadmap

• CNN architectures (ResNet, EfficientNet)
• Vision Transformers support
• Multi-label classification
• Object detection capabilities
• TensorRT optimization
• ONNX export support
• Gradio web interface
• Docker containerization

---

⭐ If you find this project helpful, please give it a star on GitHub!

Made with ❤️ by the Universal Image Classifier team