MycoNom / app.py
JohanBeytell's picture
Update app.py
a8e35e5 verified
import gradio as gr
import joblib
import pandas as pd
# Load model and mappings
model = joblib.load('mushroom_classifier.pkl')
mappings = joblib.load('mappings.pkl')
feature_options = {
'cap-shape': {'b': 'bell', 'c': 'conical', 'x': 'convex', 'f': 'flat', 'k': 'knobbed', 's': 'sunken'},
'cap-surface': {'f': 'fibrous', 'g': 'grooves', 'y': 'scaly', 's': 'smooth'},
'cap-color': {'n': 'brown', 'b': 'buff', 'c': 'cinnamon', 'g': 'gray', 'r': 'green', 'p': 'pink', 'u': 'purple', 'e': 'red', 'w': 'white', 'y': 'yellow'},
'bruises': {'t': 'bruises', 'f': 'no'},
'odor': {'a': 'almond', 'l': 'anise', 'c': 'creosote', 'y': 'fishy', 'f': 'foul', 'm': 'musty', 'n': 'none', 'p': 'pungent', 's': 'spicy'},
'gill-attachment': {'a': 'attached', 'd': 'descending', 'f': 'free', 'n': 'notched'},
'gill-spacing': {'c': 'close', 'w': 'crowded', 'd': 'distant'},
'gill-size': {'b': 'broad', 'n': 'narrow'},
'gill-color': {'k': 'black', 'n': 'brown', 'b': 'buff', 'h': 'chocolate', 'g': 'gray', 'r': 'green', 'o': 'orange', 'p': 'pink', 'u': 'purple', 'e': 'red', 'w': 'white', 'y': 'yellow'},
'stalk-shape': {'e': 'enlarging', 't': 'tapering'},
'stalk-root': {'b': 'bulbous', 'c': 'club', 'u': 'cup', 'e': 'equal', 'z': 'rhizomorphs', 'r': 'rooted', '?': 'missing'},
'stalk-surface-above-ring': {'f': 'fibrous', 'y': 'scaly', 'k': 'silky', 's': 'smooth'},
'stalk-surface-below-ring': {'f': 'fibrous', 'y': 'scaly', 'k': 'silky', 's': 'smooth'},
'stalk-color-above-ring': {'n': 'brown', 'b': 'buff', 'c': 'cinnamon', 'g': 'gray', 'o': 'orange', 'p': 'pink', 'e': 'red', 'w': 'white', 'y': 'yellow'},
'stalk-color-below-ring': {'n': 'brown', 'b': 'buff', 'c': 'cinnamon', 'g': 'gray', 'o': 'orange', 'p': 'pink', 'e': 'red', 'w': 'white', 'y': 'yellow'},
'veil-type': {'p': 'partial', 'u': 'universal'},
'veil-color': {'n': 'brown', 'o': 'orange', 'w': 'white', 'y': 'yellow'},
'ring-number': {'n': 'none', 'o': 'one', 't': 'two'},
'ring-type': {'c': 'cobwebby', 'e': 'evanescent', 'f': 'flaring', 'l': 'large', 'n': 'none', 'p': 'pendant', 's': 'sheathing', 'z': 'zone'},
'spore-print-color': {'k': 'black', 'n': 'brown', 'b': 'buff', 'h': 'chocolate', 'r': 'green', 'o': 'orange', 'u': 'purple', 'w': 'white', 'y': 'yellow'},
'population': {'a': 'abundant', 'c': 'clustered', 'n': 'numerous', 's': 'scattered', 'v': 'several', 'y': 'solitary'},
'habitat': {'g': 'grasses', 'l': 'leaves', 'm': 'meadows', 'p': 'paths', 'u': 'urban', 'w': 'waste', 'd': 'woods'}
}
# def predict_mushroom(features):
# numerical_features = {feature: feature_options[feature][value] for feature, value in features.items()}
# input_df = pd.DataFrame([numerical_features])
# prediction = model.predict(input_df)
# return 'Poisonous' if prediction[0] == 1 else 'Edible'
def predict_mushroom(*inputs):
# Build a dictionary of feature: full name selections.
# The keys are the feature names (e.g., 'cap-shape') and the values are the full names (e.g., 'bell')
user_full = dict(zip(feature_options.keys(), inputs))
# Convert the full names to letters using feature_options (reverse mapping)
user_letters = {}
for feature, full_value in user_full.items():
# Create reverse mapping: full name -> letter for this feature
inverse = {v: k for k, v in feature_options[feature].items()}
if full_value in inverse:
user_letters[feature] = inverse[full_value]
else:
raise ValueError(f"Invalid selection for {feature}: {full_value}")
# Convert letters to numeric values using the mappings from mappings.pkl
numeric_features = {}
for feature, letter in user_letters.items():
if feature in mappings and letter in mappings[feature]:
numeric_features[feature] = mappings[feature][letter]
else:
raise ValueError(f"Mapping not found for feature {feature} with letter {letter}")
# Create a DataFrame for model input
input_df = pd.DataFrame([numeric_features])
# Load the trained model (assumes the file is in the working directory)
model = joblib.load('mushroom_classifier.pkl')
# Predict using the model
prediction = model.predict(input_df)
# Return the human-readable result
return 'Edible' if prediction[0] == 0 else 'Poisonous'
# Example 1: Death cap mushroom (poisonous)
death_cap_example = [
"convex", # cap-shape: letter 'x'
"smooth", # cap-surface: letter 's'
"green", # cap-color: letter 'r'
"no", # bruises: letter 'f'
"almond", # odor: letter 'a' (Amanita phalloides is noted for a slight almond odor)
"free", # gill-attachment: letter 'f'
"close", # gill-spacing: letter 'c'
"broad", # gill-size: letter 'b'
"white", # gill-color: letter 'w'
"tapering", # stalk-shape: letter 't'
"bulbous", # stalk-root: letter 'b'
"smooth", # stalk-surface-above-ring: letter 's'
"smooth", # stalk-surface-below-ring: letter 's'
"white", # stalk-color-above-ring: letter 'w'
"white", # stalk-color-below-ring: letter 'w'
"universal", # veil-type: letter 'u'
"white", # veil-color: letter 'w'
"one", # ring-number: letter 'o'
"evanescent", # ring-type: letter 'e'
"white", # spore-print-color: letter 'w'
"scattered", # population: letter 's'
"woods" # habitat: letter 'd'
]
# Example 2: Normal brown edible mushroom
edible_example = [
"bell", # cap-shape ('b')
"fibrous", # cap-surface ('f')
"brown", # cap-color ('n')
"no", # bruises ('f')
"almond", # odor ('a')
"attached", # gill-attachment ('a')
"close", # gill-spacing ('c')
"broad", # gill-size ('b')
"brown", # gill-color ('n')
"enlarging", # stalk-shape ('e')
"bulbous", # stalk-root ('b')
"fibrous", # stalk-surface-above-ring ('f')
"fibrous", # stalk-surface-below-ring ('f')
"brown", # stalk-color-above-ring ('n')
"brown", # stalk-color-below-ring ('n')
"partial", # veil-type ('p')
"brown", # veil-color ('n')
"none", # ring-number ('n')
"none", # ring-type ('n')
"brown", # spore-print-color ('n')
"numerous", # population ('n')
"meadows" # habitat ('m')
]
demo = gr.Interface(
fn=predict_mushroom,
inputs=[gr.Dropdown(choices=list(options.values()), label=feature) for feature, options in feature_options.items()],
outputs=gr.Textbox(label="Prediction"),
title="MycoNom - Mushroom Edibility Classifier",
description="Select the mushroom features to determine if it's edible or poisonous.<br><br>You can train your own version of this model by heading to OPEN-ARC: https://github.com/Infinitode/OPEN-ARC.<br><br>**Disclaimer:** This model is for **educational purposes only** and should not be used for real-life mushroom classification or any decision-making processes related to the consumption of mushrooms. While the model performs well on the provided dataset, it has not been thoroughly validated for real-world scenarios and may not accurately detect poisonous mushrooms in all conditions. Always consult an expert or use trusted resources when identifying mushrooms.",
examples=[death_cap_example, edible_example]
)
demo.launch()