Views
No views yet
1import torch
2from PIL import Image
3from torchvision import models, transforms
4from huggingface_hub import hf_hub_download
5
6ckpt = torch.load(
7 hf_hub_download("delcenjo/flower-image-classifier", "flower_classifier.pt"),
8 map_location="cpu",
9)
10classes = ckpt["classes"]
11
12model = models.resnet18(weights=None)
13model.fc = torch.nn.Linear(model.fc.in_features, len(classes))
14model.load_state_dict(ckpt["model_state"])
15model.eval()
16
17preprocess = transforms.Compose([
18 transforms.Resize((128, 128)),
19 transforms.ToTensor(),
20 transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
21])
22
23image = Image.open("flower.jpg").convert("RGB")
24with torch.no_grad():
25 probs = model(preprocess(image).unsqueeze(0)).softmax(dim=1)[0]
26print(classes[int(probs.argmax())], float(probs.max()))