Arko007 / walnut-defect-classifier

Walnut Shell Defect Classifier

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.

Classes

Metrics (Epoch 8 — Best Checkpoint)

Val Accuracy: 98.55% | Macro F1: 0.98

Training Setup

Inference

import torch, timm
from PIL import Image
import torchvision.transforms as transforms

CLASSES = ["Healthy", "Black Spot", "Shriveled", "Damaged"]

model = timm.create_model("efficientnet_b3", pretrained=False,
                           num_classes=4, drop_rate=0.4)
ckpt  = torch.load("best_model.pth", map_location="cpu")
state = {k.replace("module.", ""): v for k, v in ckpt["model_state_dict"].items()}
model.load_state_dict(state)
model.eval()

transform = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

img    = Image.open("walnut.jpg").convert("RGB")
x      = transform(img).unsqueeze(0)
probs  = torch.softmax(model(x), dim=1)
conf, idx = probs.max(0)
print({"defect_class": CLASSES[idx.item()], "confidence": round(conf.item(), 4)})

License

Apache 2.0