EfficientNet-B3 finetuned for walnut shell defect classification across 4 categories.
Trained on the Nut Surface Defect Dataset with class remapping to match walnut-specific defect taxonomy.
1import torch, timm
2from PIL import Image
3import torchvision.transforms as transforms
4
5CLASSES = ["Healthy", "Black Spot", "Shriveled", "Damaged"]
6
7model = timm.create_model("efficientnet_b3", pretrained=False,
8 num_classes=4, drop_rate=0.4)
9ckpt = torch.load("best_model.pth", map_location="cpu")
10state = {k.replace("module.", ""): v for k, v in ckpt["model_state_dict"].items()}
11model.load_state_dict(state)
12model.eval()
13
14transform = transforms.Compose([
15 transforms.Resize((512, 512)),
16 transforms.ToTensor(),
17 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
18])
19
20img = Image.open("walnut.jpg").convert("RGB")
21x = transform(img).unsqueeze(0)
22probs = torch.softmax(model(x), dim=1)
23conf, idx = probs.max(0)
24print({"defect_class": CLASSES[idx.item()], "confidence": round(conf.item(), 4)})
