|
|
---
|
|
|
license: mit
|
|
|
tags:
|
|
|
- image-classification
|
|
|
- bacteria
|
|
|
- medical-imaging
|
|
|
- efficientnet
|
|
|
- pytorch
|
|
|
- dibas
|
|
|
datasets:
|
|
|
- custom
|
|
|
metrics:
|
|
|
- accuracy
|
|
|
- f1
|
|
|
library_name: timm
|
|
|
pipeline_tag: image-classification
|
|
|
---
|
|
|
|
|
|
# EfficientNet-B0 for Bacterial Colony Classification
|
|
|
|
|
|
This model is a fine-tuned version of **EfficientNet-B0** on the [DIBaS (Digital Image of Bacterial Species)](http://misztal.edu.pl/software/databases/dibas/) dataset for classifying bacterial colony images into 33 species.
|
|
|
|
|
|
## Model Description
|
|
|
|
|
|
- **Model Architecture:** EfficientNet-B0 (pretrained on ImageNet)
|
|
|
- **Task:** Multi-class image classification (33 bacterial species)
|
|
|
- **Dataset:** DIBaS - 660+ microscopy images of bacterial colonies
|
|
|
- **Framework:** PyTorch + timm
|
|
|
|
|
|
## Performance
|
|
|
|
|
|
| Metric | Value |
|
|
|
|--------|-------|
|
|
|
| **Validation Accuracy** | 91.67% |
|
|
|
| **Macro F1-Score** | 0.917 |
|
|
|
| **Parameters** | 4.05M |
|
|
|
| **Model Size** | 15.7 MB |
|
|
|
| **GPU Latency** | 5.81 ms (RTX 4070 SUPER) |
|
|
|
| **CPU Latency** | 25.76 ms |
|
|
|
|
|
|
### Comparison with Other Models
|
|
|
|
|
|
| Model | Params (M) | Val Accuracy |
|
|
|
|-------|------------|--------------|
|
|
|
| MobileNetV3-Large | 4.24 | **95.45%** |
|
|
|
| ResNet50 | 23.58 | 93.94% |
|
|
|
| **EfficientNet-B0** | 4.05 | 91.67% |
|
|
|
|
|
|
## Why Choose EfficientNet-B0?
|
|
|
|
|
|
- ✅ **Smallest parameter count** (4.05M) among top performers
|
|
|
- ✅ **Excellent accuracy-to-size ratio**
|
|
|
- ✅ **Good balance** between speed and accuracy
|
|
|
- ✅ **Mobile-friendly** for edge deployment
|
|
|
|
|
|
## Training Details
|
|
|
|
|
|
- **Optimizer:** AdamW (lr=1e-3, weight_decay=1e-4)
|
|
|
- **Epochs:** 10 (early convergence)
|
|
|
- **Batch Size:** 32
|
|
|
- **Image Size:** 224×224
|
|
|
- **Augmentation:** RandomResizedCrop, HorizontalFlip
|
|
|
- **Hardware:** NVIDIA RTX 4070 SUPER
|
|
|
- **Mixed Precision:** Enabled (AMP)
|
|
|
- **Train/Val/Test Split:** 70/20/10 (stratified, seed=42)
|
|
|
|
|
|
## How to Use
|
|
|
|
|
|
```python
|
|
|
import timm
|
|
|
import torch
|
|
|
from PIL import Image
|
|
|
from torchvision import transforms
|
|
|
|
|
|
# Load model
|
|
|
model = timm.create_model('efficientnet_b0', pretrained=False, num_classes=33)
|
|
|
state_dict = torch.load('pytorch_model.bin', map_location='cpu')
|
|
|
model.load_state_dict(state_dict)
|
|
|
model.eval()
|
|
|
|
|
|
# Preprocessing
|
|
|
transform = transforms.Compose([
|
|
|
transforms.Resize(256),
|
|
|
transforms.CenterCrop(224),
|
|
|
transforms.ToTensor(),
|
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
|
|
])
|
|
|
|
|
|
# Inference
|
|
|
image = Image.open('bacteria_image.jpg').convert('RGB')
|
|
|
input_tensor = transform(image).unsqueeze(0)
|
|
|
|
|
|
with torch.no_grad():
|
|
|
outputs = model(input_tensor)
|
|
|
predicted_class = outputs.argmax(dim=1).item()
|
|
|
|
|
|
print(f"Predicted class: {CLASS_NAMES[predicted_class]}")
|
|
|
```
|
|
|
|
|
|
### Class Labels
|
|
|
|
|
|
```python
|
|
|
CLASS_NAMES = [
|
|
|
"Acinetobacter_baumannii", "Actinomyces_israelii", "Bacteroides_fragilis",
|
|
|
"Bifidobacterium_spp", "Candida_albicans", "Clostridium_perfringens",
|
|
|
"Enterococcus_faecalis", "Enterococcus_faecium", "Escherichia_coli",
|
|
|
"Fusobacterium", "Lactobacillus_casei", "Lactobacillus_crispatus",
|
|
|
"Lactobacillus_delbrueckii", "Lactobacillus_gasseri", "Lactobacillus_jensenii",
|
|
|
"Lactobacillus_johnsonii", "Lactobacillus_paracasei", "Lactobacillus_plantarum",
|
|
|
"Lactobacillus_reuteri", "Lactobacillus_rhamnosus", "Lactobacillus_salivarius",
|
|
|
"Listeria_monocytogenes", "Micrococcus_spp", "Neisseria_gonorrhoeae",
|
|
|
"Porphyromonas_gingivalis", "Propionibacterium_acnes", "Proteus",
|
|
|
"Pseudomonas_aeruginosa", "Staphylococcus_aureus", "Staphylococcus_epidermidis",
|
|
|
"Staphylococcus_saprophyticus", "Streptococcus_agalactiae", "Veillonella"
|
|
|
]
|
|
|
```
|
|
|
|
|
|
## Limitations
|
|
|
|
|
|
- Trained on single laboratory/microscope setup (DIBaS dataset)
|
|
|
- May not generalize to different imaging conditions
|
|
|
- Not validated for clinical diagnostic use
|
|
|
|
|
|
## Related Models
|
|
|
|
|
|
- [MobileNetV3-Large](https://huggingface.co/ihoflaz/dibas-mobilenet-v3-large) - Best accuracy (95.45%)
|
|
|
- [ResNet50](https://huggingface.co/ihoflaz/dibas-resnet50) - Classic architecture (93.94%)
|
|
|
|
|
|
## Citation
|
|
|
|
|
|
```bibtex
|
|
|
@inproceedings{hoflaz2025bacterial,
|
|
|
title={Lightweight CNNs Outperform Vision Transformers for Bacterial Colony Classification},
|
|
|
author={Hoflaz, Ibrahim},
|
|
|
booktitle={IEEE Conference},
|
|
|
year={2025}
|
|
|
}
|
|
|
```
|
|
|
|
|
|
## Resources
|
|
|
|
|
|
- **GitHub:** [ihoflaz/bacterial-colony-classification](https://github.com/ihoflaz/bacterial-colony-classification)
|
|
|
- **DIBaS Dataset:** [http://misztal.edu.pl/software/databases/dibas/](http://misztal.edu.pl/software/databases/dibas/)
|
|
|
|
