food_classify / app.py
JanmejayM46's picture
Create app.py
44e0053
import streamlit as st
from transformers import pipeline
from rembg import remove
from PIL import Image
from io import BytesIO
import base64
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import random
st.set_page_config(layout="wide", page_title="Image Classification App")
st.write("## Image Food Classification App")
st.sidebar.write("## Upload and download :gear:")
# Initialize image classification and recipe generation models
image_classifier = pipeline("image-classification", model="mjsp/sweet")
recipe_model = GPT2LMHeadModel.from_pretrained("gpt2")
recipe_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
MAX_FILE_SIZE = 5 * 1024 * 1024 # 5MB
def convert_image(img):
buf = BytesIO()
img.save(buf, format="PNG")
byte_im = buf.getvalue()
return byte_im
def fix_image(upload):
image = Image.open(upload)
col1.write("Original Image :camera:")
col1.image(image)
fixed = remove(image)
col2.write("Fixed Image :wrench:")
col2.image(fixed)
st.sidebar.markdown("\n")
st.sidebar.download_button("Download fixed image", convert_image(fixed), "fixed.png", "image/png")
def generate_recipe(title, max_length=200):
matching_entries = [entry for entry in dataset if entry["title"] == title]
if matching_entries:
selected_entry = random.choice(matching_entries)
title = selected_entry["title"]
ingredients = selected_entry["ingredients"]
else:
title = "Default Recipe Title"
ingredients = []
input_text = f"Title: {title}\nIngredients: {', '.join(ingredients)}\n Instructions:"
input_ids = recipe_tokenizer.encode(input_text, return_tensors="pt")
output = recipe_model.generate(input_ids, max_length=max_length, num_return_sequences=1)
generated_recipe = recipe_tokenizer.decode(output[0], skip_special_tokens=True)
return generated_recipe
# Your dataset of titles, ingredients, and recipes
dataset = [
{
"title": "Gulab Jamun",
"ingredients": ["milk powder", "ghee", "rose water", "saffron", "cardamom", "sugar syrup"],
"recipe": "Instructions for making Gulab Jamun..."
},
{
"title": "Jalebi",
"ingredients": ["all-purpose flour", "yogurt", "sugar", "water", "saffron strands", "cardamom powder", "ghee or oil for frying"],
"recipe": "Instructions for making Jalebi..."
},
{
"title": "Rasgulla",
"ingredients": ["milk", "sugar", "lemon juice", "rose water"],
"recipe": "Instructions for making Rasgulla..."
}
]
col1, col2 = st.columns(2)
my_upload = st.sidebar.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
st.write()
st.markdown("""---""")
st.write()
# global label
# label = ""
if my_upload is not None:
st.image(my_upload, caption="Uploaded Image", use_column_width=True)
if st.sidebar.button("Classify"):
st.sidebar.text("Classifying...")
image = Image.open(my_upload)
try:
classification_result = image_classifier(image)
top_prediction = classification_result[0]
label = top_prediction['label']
score = top_prediction['score']
st.sidebar.text("Top Prediction:")
st.sidebar.text(f"Label: {label}, Score: {score:.3f}")
except Exception as e:
st.error(f"Error during classification: {str(e)}")
if my_upload.size > MAX_FILE_SIZE:
st.error("The uploaded file is too large. Please upload an image smaller than 5MB.")
else:
fix_image(upload=my_upload)
else:
fix_image("jalebi.jpg")
# Recipe generation based on selected item
st.write("## Recipe Generation")
selected_item = st.selectbox("Select a food item", [entry["title"] for entry in dataset])
if st.button("Generate Recipe"):
generated_recipe = generate_recipe(selected_item, max_length=200)
st.write(f"Recipe for {selected_item}:\n{generated_recipe}")
# Add some descriptions and instructions
st.sidebar.markdown("### Instructions")
st.sidebar.markdown("1. Upload an image.")
st.sidebar.markdown("2. Click the 'Classify' button to get the classification results.")
st.sidebar.markdown("3. Select a food item to generate a recipe.")
st.sidebar.markdown("4. Click the 'Generate Recipe' button to get the recipe.")
# Display a placeholder for the main content
st.write("Please upload an image and use the sidebar to classify it and generate a recipe.")