README.md

---
tags:
  - image-classification
  - bacteria
  - microscopy
  - transfer-learning
  - efficientnet
library_name: timm
datasets:
  - custom
pipeline_tag: image-classification
---

# Bacteria Image Classifier (EfficientNet-B0)

A fine-tuned EfficientNet-B0 model for classifying microscopy images of bacteria
and fungi into 20 classes (10 organisms × 2 imaging types: gram stain and media plate).

## Classes

The model distinguishes between the following 20 classes:

| Organism | Gram Stain | Media Plate |
|----------|-----------|-------------|
| Aspergillus niger | ✓ | ✓ |
| Bacillus subtilis | ✓ | ✓ |
| Candida albicans | ✓ | ✓ |
| Clostridium sporogenes | ✓ | ✓ |
| Enterococcus faecalis | ✓ | ✓ |
| Escherichia coli | ✓ | ✓ |
| Klebsiella pneumoniae | ✓ | ✓ |
| Pseudomonas aeruginosa | ✓ | ✓ |
| Staphylococcus aureus | ✓ | ✓ |
| Streptococcus pyogenes | ✓ | ✓ |

## Usage

```python
import timm, torch, json
from torchvision import transforms
from PIL import Image
from huggingface_hub import hf_hub_download

repo_id = "lederyou/bacteria-classifier"
model_path = hf_hub_download(repo_id, "model.pth")
class_names = json.load(open(hf_hub_download(repo_id, "class_names.json")))

model = timm.create_model("efficientnet_b0", pretrained=False, num_classes=20)
model.load_state_dict(torch.load(model_path, map_location="cpu"))
model.eval()

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

img = transform(Image.open("your_image.png").convert("RGB")).unsqueeze(0)
with torch.no_grad():
    pred = model(img).argmax(1).item()
print(class_names[pred])
```

## Training

- **Base model**: EfficientNet-B0 (pretrained on ImageNet)
- **Method**: Two-phase transfer learning (frozen backbone → full fine-tuning)
- **Dataset**: 629 images, 20 classes, 70/15/15 train/val/test split