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())