An EfficientNet-b3 model fine-tuned for classifying mushroom images into the 12 popular Russian mushroom classes corresponding to species and edibility. Fine-tuned on the
12_popular_russia_mushrooms_edible_poisonous dataset.
Here is how to use Musheff model to classify a mushroom image into one of the 12 Russian classes (species and edibility):
1import random
2
3import torch
4
5from datasets import load_dataset
6from transformers import (
7 AutoImageProcessor,
8 AutoModel,
9)
10
11test_dataset = load_dataset(
12 "SoFa325/12_popular_russia_mushrooms_edible_poisonous", split="test"
13)
14
15test_len = len(test_dataset)
16
17# Pick a random image from test set
18random_index = random.randint(0, test_len)
19
20image = test_dataset["image"][random_index]
21
22preprocessor = AutoImageProcessor.from_pretrained(
23 "blasisd/musheff",
24 trust_remote_code=True,
25 use_fast=True,
26)
27
28
29model = AutoModel.from_pretrained(
30 "blasisd/musheff",
31 trust_remote_code=True,
32 low_cpu_mem_usage=True, # Activates memory-efficient loading
33 device_map="auto", # Distributes layers across devices
34)
35
36inputs = preprocessor(image, return_tensors="pt").to(model.device)
37
38model.eval()
39with torch.inference_mode():
40 logits = model(inputs["pixel_values"])
41
42# model predicts one of the 12 potential mushroom classes
43predicted_label = logits.argmax(dim=1).item()
44
45print(f"True label: {test_dataset['label'][random_index]}")
46print(f"Predicted label: {model.config.id2label[predicted_label]}"),
Alternatively, you can download the files from the
repository locally
and follow the steps below:
1import json
2import random
3
4import torch
5
6from datasets import load_dataset
7from torchvision import models
8
9from model import Musheff
10
11
12# Device agnostic
13device = "cuda" if torch.cuda.is_available() else "cpu"
14
15test_dataset = load_dataset(
16 "SoFa325/12_popular_russia_mushrooms_edible_poisonous", split="test"
17)
18
19test_len = len(test_dataset)
20
21# Pick a random image from test set
22random_index = random.randint(0, test_len)
23
24image = test_dataset["image"][random_index]
25
26with open("config.json", "r") as json_fp:
27 config = json.load(json_fp)
28
29model = Musheff(config)
30model.model.load_state_dict(torch.load("musheff.pth"))
31model.to(device=device)
32
33transform = models.EfficientNet_B3_Weights.DEFAULT.transforms()
34
35img = transform(image)
36
37# Expecting 4D shape i.e. (batch_size, channels, height, width)
38img = img.unsqueeze(0)
39
40model.eval()
41
42with torch.inference_mode():
43 logits = model(img.to(device))
44
45# model predicts one of the 12 potential mushroom classes
46predicted_label = logits.argmax(dim=1).item()
47
48with open("config.json", "r") as json_fp:
49 id2label = json.load(json_fp).get("id2label")
50
51print(f"True label: {test_dataset['label'][random_index]}")
52print(f"Predicted label: {id2label[str(predicted_label)]}")
Install required Python packages using either method.
Explore fine-tuned variants on the
Hugging Face Hub