Bwenge840/vit-base-patch16-384-coffee-preloaded

Vision Transformer coffee image classifier fine-tuned with vit_base_patch16_384 at 384x384 resolution.

Files

• vit_base_patch16_384_coffee_preloaded.pth: PyTorch checkpoint.

This repository also includes ONNX and TFLite exports when uploaded with --include-exports.

Expected Input Specifications

• Channels: RGB
• Image Size: 384x384 (e.g. (384, 384) with bicubic interpolation)
• Normalization: Mean [0.5, 0.5, 0.5], std [0.5, 0.5, 0.5]

Classes

• KR1
• KR10
• KR3
• KR4
• KR5
• KR6
• KR7
• KR8
• KR9

Python Inference Example

import timm
import torch
import torch.nn as nn
from PIL import Image
import torchvision.transforms as T
from huggingface_hub import hf_hub_download

# Define model architecture matching fastai training wrapper
class PyTorchViTModel(nn.Module):
    def __init__(self, num_classes=9):
        super().__init__()
        self.backbone = timm.create_model("vit_base_patch16_384", pretrained=False, num_classes=0)
        self.head = nn.Sequential(
            nn.BatchNorm1d(768),
            nn.Dropout(p=0.25),
            nn.Linear(768, 512, bias=False),
            nn.ReLU(inplace=True),
            nn.BatchNorm1d(512),
            nn.Dropout(p=0.5),
            nn.Linear(512, num_classes, bias=False)
        )
        
    def forward(self, x):
        return self.head(self.backbone(x))

# 1. Load weights
model = PyTorchViTModel(num_classes=9)
weights_path = hf_hub_download(repo_id="Bwenge840/vit-base-patch16-384-coffee-preloaded", filename="vit_base_patch16_384_coffee_preloaded.pth")
checkpoint = torch.load(weights_path, map_location="cpu", weights_only=False)
state_dict = checkpoint.get("model", checkpoint)
remapped = {
    (k.replace("0.model.", "backbone.").replace("1.", "head.") if k.startswith(("0.model.", "1.")) else k): v
    for k, v in state_dict.items()
}
model.load_state_dict(remapped)
model.eval()

# 2. Preprocess test image (384x384)
transform = T.Compose([
    T.Resize((384, 384), interpolation=T.InterpolationMode.BICUBIC),
    T.ToTensor(),
    T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])

img = Image.open("coffee_leaf.jpg").convert("RGB")
input_tensor = transform(img).unsqueeze(0)

# 3. Predict
classes = ['KR1', 'KR10', 'KR3', 'KR4', 'KR5', 'KR6', 'KR7', 'KR8', 'KR9']
with torch.no_grad():
    logits = model(input_tensor)
    probs = torch.softmax(logits, dim=1)[0]
    idx = probs.argmax().item()
    print(f"Predicted variety: {classes[idx]} (Confidence: {probs[idx]:.2%})")