This model is a fine-tuned version of
ResNet50 on the
DIBaS (Digital Image of Bacterial Species) dataset for classifying bacterial colony images into 33 species.
1import timm
2import torch
3from PIL import Image
4from torchvision import transforms
5
6# Load model
7model = timm.create_model('resnet50', pretrained=False, num_classes=33)
8state_dict = torch.load('pytorch_model.bin', map_location='cpu')
9model.load_state_dict(state_dict)
10model.eval()
11
12# Preprocessing
13transform = transforms.Compose([
14 transforms.Resize(256),
15 transforms.CenterCrop(224),
16 transforms.ToTensor(),
17 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
18])
19
20# Inference
21image = Image.open('bacteria_image.jpg').convert('RGB')
22input_tensor = transform(image).unsqueeze(0)
23
24with torch.no_grad():
25 outputs = model(input_tensor)
26 predicted_class = outputs.argmax(dim=1).item()
27
28print(f"Predicted class: {CLASS_NAMES[predicted_class]}")
1CLASS_NAMES = [
2 "Acinetobacter_baumannii", "Actinomyces_israelii", "Bacteroides_fragilis",
3 "Bifidobacterium_spp", "Candida_albicans", "Clostridium_perfringens",
4 "Enterococcus_faecalis", "Enterococcus_faecium", "Escherichia_coli",
5 "Fusobacterium", "Lactobacillus_casei", "Lactobacillus_crispatus",
6 "Lactobacillus_delbrueckii", "Lactobacillus_gasseri", "Lactobacillus_jensenii",
7 "Lactobacillus_johnsonii", "Lactobacillus_paracasei", "Lactobacillus_plantarum",
8 "Lactobacillus_reuteri", "Lactobacillus_rhamnosus", "Lactobacillus_salivarius",
9 "Listeria_monocytogenes", "Micrococcus_spp", "Neisseria_gonorrhoeae",
10 "Porphyromonas_gingivalis", "Propionibacterium_acnes", "Proteus",
11 "Pseudomonas_aeruginosa", "Staphylococcus_aureus", "Staphylococcus_epidermidis",
12 "Staphylococcus_saprophyticus", "Streptococcus_agalactiae", "Veillonella"
13]
1@inproceedings{hoflaz2025bacterial,
2 title={Lightweight CNNs Outperform Vision Transformers for Bacterial Colony Classification},
3 author={Hoflaz, Ibrahim},
4 booktitle={IEEE Conference},
5 year={2025}
6}
