lederyou/bacteria-classifier

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

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