Spaces:
Sleeping
Sleeping
File size: 4,183 Bytes
76ebb93 8176067 5d3c722 a9fec20 8176067 4f0dcd5 a9fec20 4f0dcd5 8176067 4f0dcd5 8176067 4f0dcd5 a9fec20 8176067 6cd8801 fbcff5b 8176067 fbcff5b 76ebb93 5d3c722 76ebb93 8176067 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
# streamlit_app.py
import streamlit as st
from fastai.vision.all import *
import shutil
import os
# Function to get the label from the file name
def get_label(file_name):
return file_name.split('-')[0]
# Function to prepare data (similar to your code)
def prepare_data(food_path, label_a, label_b):
for img in get_image_files(food_path):
if label_a in str(img):
img.rename(f"{img.parent}/{label_a}-{img.name}")
elif label_b in str(img):
img.rename(f"{img.parent}/{label_b}-{img.name}")
else:
os.remove(img)
# Function to train the model
def train_model(food_path, label_func):
dls = ImageDataLoaders.from_name_func(
food_path, get_image_files(food_path), valid_pct=0.2, seed=420,
label_func=label_func, item_tfms=Resize(230)
)
learn = cnn_learner(dls, resnet34, metrics=error_rate, pretrained=True)
learn.fine_tune(epochs=1)
return learn
# ... (previous code)
# ... (previous code)
# Streamlit app
def main():
st.title("Food Classifier Streamlit App")
# Sidebar options
options = ["Train Model", "Upload Image", "Test Random Images", "Confusion Matrix"]
choice = st.sidebar.selectbox("Choose an option", options)
if choice == "Train Model":
st.subheader("Training the Model")
food_path = Path("~/.fastai/data/food-101/food-101").expanduser()
if not food_path.exists():
try:
food_path = untar_data(URLs.FOOD)
except FileExistsError:
st.warning("Data directory already exists. Skipping download.")
label_a = st.text_input("Enter label A:", "samosa")
label_b = st.text_input("Enter label B:", "hot_and_sour_soup")
prepare_data(food_path, label_a, label_b)
learn = train_model(food_path, get_label)
st.session_state.model = learn # Save the model to session state
st.success("Model trained successfully!")
# ... (rest of the code remains unchanged)
elif choice == "Upload Image":
st.subheader("Upload Your Own Images")
if "model" not in st.session_state:
st.warning("Please train the model first.")
else:
uploaded_files = st.file_uploader("Choose images", type=["jpg", "jpeg", "png"], accept_multiple_files=True)
if uploaded_files:
for img in uploaded_files:
img = PILImage.create(img)
label, _, probs = st.session_state.model.predict(img)
st.image(img, caption=f"This is a {label}.")
st.write(f"{label_a}: {probs[1].item():.6f}")
st.write(f"{label_b}: {probs[0].item():.6f}")
if uploaded_files:
for img in uploaded_files:
img = PILImage.create(img)
label, _, probs = st.session_state.model.predict(img)
st.image(img, caption=f"This is a {label}.")
st.write(f"{label_a}: {probs[1].item():.6f}")
st.write(f"{label_b}: {probs[0].item():.6f}")
elif choice == "Test Random Images":
st.subheader("Test Using Images in Dataset")
if "model" not in st.session_state:
st.warning("Please train the model first.")
else:
for i in range(0, 5): # Change 5 to the number of images you want to display
random_index = random.randint(0, len(get_image_files(food_path)) - 1)
img_path = get_image_files(food_path)[random_index]
img = mpimg.imread(img_path)
label, _, probs = st.session_state.model.predict(img)
st.image(img, caption=f"Predicted label: {label}")
elif choice == "Confusion Matrix":
st.subheader("Confusion Matrix")
if "model" not in st.session_state:
st.warning("Please train the model first.")
else:
interp = ClassificationInterpretation.from_learner(st.session_state.model)
st.pyplot(interp.plot_confusion_matrix())
# Run the Streamlit app
if __name__ == "__main__":
main()
|