Spaces:
Sleeping
Sleeping
import streamlit as st | |
import os | |
import json | |
import cv2 | |
import pandas as pd | |
from ultralytics import YOLO | |
from PIL import Image | |
# Configure Streamlit pages | |
st.set_page_config( | |
page_title="NutriDetect: AI-powered Food Detection and Nutritional Insights", | |
page_icon="π΄", | |
# Set layout to wide for better visualization | |
) | |
# Sidebar navigation | |
st.sidebar.title("Navigation") | |
page = st.sidebar.radio("Go to", ["Home", "Description", "Try Models"]) | |
# Load nutritional data from the JSON file | |
def load_nutritional_data(json_path): | |
with open(json_path, 'r') as file: | |
data = json.load(file) | |
for food in data["food"]: | |
food["preprocessed_name"] = food["name"].lower().replace(",", "").strip() | |
return data | |
# Load models | |
def load_models(): | |
yolov8_model = YOLO("./model/Yolov8m.pt") | |
yolov11_model = YOLO("./model/Yolo11m.pt") | |
return yolov8_model, yolov11_model | |
# Global variables for models and nutritional data | |
models = load_models() | |
nutritional_json_path = "db.json" | |
nutritional_data = load_nutritional_data(nutritional_json_path) | |
# Home Page | |
if page == "Home": | |
st.title("π΄ NutriDetect: Food Detection and Nutritional Insights") | |
st.markdown("**Welcome to the NutriDetect App!** Leverage cutting-edge AI to identify food items and provide nutritional details.") | |
st.image("back_image.png", use_container_width=True, caption="AI-powered Food Detection") | |
st.subheader("π Key Features") | |
st.markdown(""" | |
- **Accurate Food Detection**: Identify food items in images. | |
- **Nutritional Insights**: Get calorie, protein, fat, and other details for detected food. | |
""") | |
# Description Page | |
elif page == "Description": | |
st.title("π About the NutriDetect App") | |
st.markdown(""" | |
The NutriDetect App uses YOLO-based models for accurate food detection and provides: | |
- Bounding boxes for detected food items. | |
- Detailed nutritional insights for the detected items. | |
""") | |
# Try Models Page | |
elif page == "Try Models": | |
st.title("π― Try the Models") | |
uploaded_file = st.file_uploader("Upload an image (jpg, jpeg, png)", type=["jpg", "jpeg", "png"]) | |
if uploaded_file: | |
st.subheader("Step 2: Choose Model") | |
model_option = st.selectbox("Select Model", ["YOLOv8m", "YOLO11m"]) | |
if st.button("π Detect"): | |
with st.spinner("Processing..."): | |
# Save the uploaded file temporarily | |
temp_file_path = "temp_image.jpg" | |
with open(temp_file_path, "wb") as f: | |
f.write(uploaded_file.getbuffer()) | |
# Perform detection | |
model = models[0] if model_option == "YOLOv8m" else models[1] | |
results = model(temp_file_path) | |
# Process detected items | |
detected_items = [] | |
for box in results[0].boxes.data.cpu().numpy(): | |
label = int(box[5]) # Get the class label | |
class_name = model.names[label] | |
nutrition_info = next( | |
(food["nutrients"] for food in nutritional_data["food"] if food["preprocessed_name"] == class_name.lower().replace(",", "").strip()), | |
None | |
) | |
if nutrition_info: | |
detected_items.append( | |
{ | |
"Food": class_name, | |
"Calories": round(nutrition_info.get("calories", 0), 2), | |
"Protein": round(nutrition_info.get("protein", 0), 2), | |
"Fat": round(nutrition_info.get("fat", 0), 2), | |
"Carbs": round(nutrition_info.get("carbs", 0), 2), | |
"Fiber": round(nutrition_info.get("fiber", 0), 2), | |
} | |
) | |
# Handle case with no detected items | |
if not detected_items: | |
st.warning("No food items detected, sorry!") | |
else: | |
# Convert detected items to DataFrame | |
df = pd.DataFrame(detected_items) | |
# Calculate totals | |
totals = { | |
"Calories": round(df["Calories"].sum(), 2), | |
"Protein": round(df["Protein"].sum(), 2), | |
"Fat": round(df["Fat"].sum(), 2), | |
"Carbs": round(df["Carbs"].sum(), 2), | |
"Fiber": round(df["Fiber"].sum(), 2), | |
} | |
# Display annotated image | |
annotated_image = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB) | |
st.image(annotated_image, caption=f"{model_option} Results", use_container_width=True) | |
# Two-column layout for detected items and totals | |
left_col, right_col = st.columns([2, 1]) | |
# Left column: Individual food items | |
with left_col: | |
st.header("Detected Food Items") | |
for _, row in df.iterrows(): | |
st.subheader(row["Food"]) | |
st.write(f"Calories: {row['Calories']} Cal") | |
st.write(f"Protein: {row['Protein']} g") | |
st.write(f"Fat: {row['Fat']} g") | |
st.write(f"Carbs: {row['Carbs']} g") | |
st.write(f"Fiber: {row['Fiber']} g") | |
st.write("---") | |
# Right column: Total nutritional summary | |
# Right column: Total nutritional summary | |
# Right column: Total nutritional summary | |
with right_col: | |
st.header("Total Nutritional Summary") | |
st.metric("Total Calories", f"{totals['Calories']} Cal") | |
st.metric("Total Protein", f"{totals['Protein']} g") | |
st.metric("Total Fat", f"{totals['Fat']} g") | |
st.metric("Total Carbs", f"{totals['Carbs']} g") | |
st.metric("Total Fiber", f"{totals['Fiber']} g") | |
# Check for unhealthy items and provide recommendations | |
detected_foods = df["Food"].str.lower().tolist() # Convert to lowercase for matching | |
unhealthy_items = [item for item in detected_foods if item in ["pizza", "hamburger", "french fries"]] | |
if unhealthy_items: | |
st.warning( | |
f"Detected unhealthy food items: {', '.join(unhealthy_items)}. " | |
"Consider replacing them with fruits and green salad for better health. π₯π" | |
) | |
else: | |
st.success("Balanced meals lead to a healthier life! π") | |