--- 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.