Image Classifier using EfficientNet

A simple image classifier using the pre-trained EfficientNet-B0 model from PyTorch.

Features

• Uses EfficientNet-B0 pre-trained model
• Supports image classification for 1000 ImageNet classes
• Fine-tune the model on your own dataset
• Load custom trained model weights
• Returns top 15 predictions with confidence scores
• Easy to use API
• Supports both CPU and GPU (if available)

Installation

1. Create a virtual environment (recommended):

   python -m venv venv
   source venv/bin/activate  # On Windows: venv\Scripts\activate

2. Install dependencies:

   pip install -r requirements.txt

Usage

1. Import and use the classifier in your code:

   from src.image_classifier.classifier import ImageClassifier
   
   # Initialize the classifier
   classifier = ImageClassifier()
   
   # Make predictions
   predictions = classifier.predict("path/to/your/image.jpg")
   
   # Print results
   for i, pred in enumerate(predictions, 1):
       print(f"{i}. {pred['class']}: {pred['confidence']}")

2. Or run the example script:

   python src/image_classifier/classifier.py

   (Remember to update the image path in the script)

Training

You can fine-tune the model on your own dataset using the command line. Your dataset should be organized into subdirectories, where each subdirectory name corresponds to a class name and contains images of that class.

python src/image_classifier/classifier.py --data_dir path/to/your/dataset --epochs 15 --batch_size 64 --learning_rate 0.001

Arguments:

• --data_dir (required): Path to the directory containing your training images, organized by class folders.
• --epochs (optional): Number of training epochs (default: 10).
• --batch_size (optional): Number of images per batch (default: 32).
• --learning_rate (optional): Learning rate for the optimizer (default: 0.001).

The trained model weights will be saved as trained_model.pth in the project root directory.

Loading a Trained Model

After training, you can load your custom model weights:

from src.image_classifier.classifier import ImageClassifier

# Initialize the classifier (loads default weights initially)
classifier = ImageClassifier()

# Load your trained weights
classifier.load_trained_model("trained_model.pth")

# Now you can use predict() with your fine-tuned model
predictions = classifier.predict("path/to/test_image.jpg")
print(predictions)

Example Output (Prediction - Default Model)

Top 5 Predictions:
1. golden retriever: 98.45%
2. Labrador retriever: 1.32%
3. German shepherd: 0.12%
4. Rottweiler: 0.08%
5. Doberman: 0.03%

Example Output (Training)

This is a sample of what the console output might look like during training:

Epoch 1, Batch 10, Loss: 1.532, Accuracy: 45.31%
Epoch 1, Batch 20, Loss: 0.987, Accuracy: 58.75%
...
Epoch 1 completed. Accuracy: 72.10%
Epoch 2, Batch 10, Loss: 0.654, Accuracy: 78.91%
...
Epoch 10 completed. Accuracy: 95.67%
Training completed and model saved.

Example Output (Prediction - Trained Model)

After loading your trained_model.pth (assuming it was trained on a dataset with 'cat' and 'dog' classes), the prediction output might look like this:

Top 5 Predictions:
1. cat: 97.21%
2. dog: 2.55%
3. tabby: 0.15%
4. tiger_cat: 0.05%
5. lynx: 0.02%

Notes

• The model expects RGB images
• Images are automatically resized and normalized
• The classifier supports common image formats (JPEG, PNG, etc.)
• GPU acceleration is automatically used if available