Huyt/cabbage-disease-rexnet150

Cabbage Disease Classifier (rexnet_150)

A rexnet_150 (timm) image classifier fine-tuned to detect cabbage plant diseases.

• Classes (8): ['Alternaria_Leaf_Spot', 'club root', 'Downy Mildew', 'Cabbage aphid colony', 'ring spot', 'Black Rot', 'Bacterial spot rot', 'No disease']
• Input size: 224x224
• Normalization: mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]

How to use

See the "Inference: How to Use This Model" section of the training notebook, or the snippet below:

import json, torch, timm
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import transforms as T

REPO_ID = "Huyt/cabbage-disease-rexnet150"

config_path = hf_hub_download(repo_id=REPO_ID, filename="config.json")
weights_path = hf_hub_download(repo_id=REPO_ID, filename="cabbage_best_model.pth")

with open(config_path) as f:
    config = json.load(f)
idx_to_class = {v: k for k, v in config["classes"].items()}

device = "cuda" if torch.cuda.is_available() else "cpu"
model = timm.create_model(config["model_name"], pretrained=False, num_classes=config["num_classes"])
model.load_state_dict(torch.load(weights_path, map_location=device))
model.to(device).eval()

tfs = T.Compose([
    T.Resize((config["image_size"], config["image_size"])),
    T.ToTensor(),
    T.Normalize(mean=config["mean"], std=config["std"]),
])

im = Image.open("path/to/image.jpg").convert("RGB")
x = tfs(im).unsqueeze(0).to(device)
with torch.no_grad():
    probs = torch.softmax(model(x), dim=1)[0]
pred_idx = int(torch.argmax(probs))
print(idx_to_class[pred_idx], probs[pred_idx].item())