| --- |
| license: mit |
| library_name: pytorch |
| pipeline_tag: image-classification |
| tags: |
| - image-classification |
| - pytorch |
| - transfer-learning |
| - resnet18 |
| - computer-vision |
| - flowers |
| --- |
| |
| # 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 |
|
|
| ```python |
| 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. |
|
|