Views
No views yet
weights=None, trained from scratch/fine-tuned)nn.Linear(num_features, len(class_names))class_names and model_state_dict keys1from huggingface_hub import hf_hub_download
2import torch, torch.nn as nn
3from torchvision import models, transforms
4from PIL import Image
5
6model_path = hf_hub_download("AAYUSHSAVALIYA/agri-ai-model", "model.pth")
7checkpoint = torch.load(model_path, map_location="cpu")
8class_names = checkpoint["class_names"]
9
10model = models.mobilenet_v2(weights=None)
11model.classifier[1] = nn.Linear(model.classifier[1].in_features, len(class_names))
12model.load_state_dict(checkpoint["model_state_dict"])
13model.eval()
14
15transform = transforms.Compose([
16 transforms.Resize((224, 224)),
17 transforms.ToTensor(),
18 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
19])
20
21img = Image.open("leaf.jpg").convert("RGB")
22with torch.no_grad():
23 output = model(transform(img).unsqueeze(0))
24 pred = output.argmax(dim=1).item()
25print(class_names[pred])class_names in the checkpoint]