CTP_Project / app.py
HassanDataSci's picture
Update app.py
37222e0 verified
raw
history blame
2.59 kB
import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from PIL import Image
# Load the image classification pipeline
@st.cache_resource
def load_image_classification_pipeline():
"""
Load the image classification pipeline using a pretrained model.
"""
return pipeline("image-classification", model="Shresthadev403/food-image-classification")
pipe_classification = load_image_classification_pipeline()
# Load Qwen tokenizer and model
@st.cache_resource
def load_qwen_model():
"""
Load the Qwen/Qwen2.5-Coder-32B-Instruct model and tokenizer.
"""
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-32B-Instruct")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-Coder-32B-Instruct", device_map="auto")
return tokenizer, model
# Function to generate ingredients using Qwen
def get_ingredients_qwen(food_name, tokenizer, model):
"""
Generate a list of ingredients for the given food item using the Qwen model.
"""
prompt = f"List the main ingredients typically used to prepare {food_name}:"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=50)
return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
# Streamlit app
st.title("Food Image Recognition with Ingredients")
# # Add the provided image as a banner
# st.image("CTP_Project/IR_IMAGE", caption="Food Recognition Model", use_column_width=True)
# Sidebar for model information
st.sidebar.title("Model Information")
st.sidebar.write("**Image Classification Model**: Shresthadev403/food-image-classification")
st.sidebar.write("**LLM for Ingredients**: Qwen2.5-Coder-32B-Instruct")
# Upload image
uploaded_file = st.file_uploader("Choose a food image...", type=["jpg", "png", "jpeg"])
# Load the Qwen model and tokenizer
tokenizer, model = load_qwen_model()
if uploaded_file is not None:
# Display the uploaded image
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Image", use_column_width=True)
st.write("Classifying...")
# Make predictions
predictions = pipe_classification(image)
# Display only the top prediction
top_food = predictions[0]['label']
st.header(f"Food: {top_food}")
# Generate and display ingredients for the top prediction
st.subheader("Ingredients")
try:
ingredients = get_ingredients_qwen(top_food, tokenizer, model)
st.write(ingredients)
except Exception as e:
st.error(f"Error generating ingredients: {e}")