delcenjo/flower-image-classifier

Flower Image Classifier (ResNet-18, transfer learning)

Try the live demo: https://huggingface.co/spaces/delcenjo/flower-classifier-demo Code on GitHub: https://github.com/delcenjo/flower-image-classifier

A small image classifier that recognises five flower species (daisy, dandelion, roses, sunflowers, tulips) using transfer learning on a pre-trained ResNet-18. The ImageNet backbone is frozen and only a new classification head is trained, so it runs well on CPU.

• Architecture: ResNet-18 (ImageNet weights), final layer replaced with a 5-class head
• Input: RGB image resized to 128x128, normalised with ImageNet statistics
• Test accuracy: about 0.77 (5 balanced classes; random baseline 0.20)
• Dataset: TensorFlow Flowers (about 3,670 images)

Usage

import torch
from PIL import Image
from torchvision import models, transforms
from huggingface_hub import hf_hub_download

ckpt = torch.load(
    hf_hub_download("delcenjo/flower-image-classifier", "flower_classifier.pt"),
    map_location="cpu",
)
classes = ckpt["classes"]

model = models.resnet18(weights=None)
model.fc = torch.nn.Linear(model.fc.in_features, len(classes))
model.load_state_dict(ckpt["model_state"])
model.eval()

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

image = Image.open("flower.jpg").convert("RGB")
with torch.no_grad():
    probs = model(preprocess(image).unsqueeze(0)).softmax(dim=1)[0]
print(classes[int(probs.argmax())], float(probs.max()))

Limitations

Trained on a small dataset at 128x128 with a frozen backbone, so accuracy is modest. Unfreezing the last ResNet block and training at 224x224 would improve it.