dibas-efficientnet-b0

EfficientNet-B0 for Bacterial Colony Classification

This model is a fine-tuned version of EfficientNet-B0 on the DIBaS (Digital Image of Bacterial Species) 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

Comparison with Other Models

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

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

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 - Best accuracy (95.45%)
• ResNet50 - Classic architecture (93.94%)

Citation

@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
• DIBaS Dataset: http://misztal.edu.pl/software/databases/dibas/