Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,7 +1,41 @@
 import streamlit as st
+import streamlit.components.v1 as components
+from transformers import pipeline
+
+
+# -------------------------
+# ADD THIS (BMI FUNCTIONS)
+# -------------------------
+def calculate_bmi(weight, height_cm):
+    height_m = height_cm / 100
+    return round(weight / (height_m ** 2), 2)
+
+def get_category(bmi):
+    if bmi < 18.5:
+        return "Underweight"
+    elif bmi < 24.9:
+        return "Normal"
+    elif bmi < 29.9:
+        return "Overweight"
+    else:
+        return "Obese"
+
+
+# -------------------------
+# YOUR MODEL (NO CHANGE)
+# -------------------------
+@st.cache_resource
+def load_model():
+    return pipeline(
+        "text-generation",
+        model="google/flan-t5-base"
+    )
+
+generator = load_model()
+
 
 # -----------------------------------
-# PAGE CONFIG
+# PAGE CONFIG (NO CHANGE)
 # -----------------------------------
 st.set_page_config(
     page_title="FitPlan AI",
@@ -9,8 +43,9 @@ st.set_page_config(
     layout="wide"
 )
 
+
 # -----------------------------------
-# CUSTOM CSS (Modern UI)
+# CUSTOM CSS (NO CHANGE)
 # -----------------------------------
 st.markdown("""
 <style>
@@ -46,12 +81,14 @@ body {
 </style>
 """, unsafe_allow_html=True)
 
+
 # -----------------------------------
-# SIDEBAR NAVIGATION
+# SIDEBAR NAVIGATION (NO CHANGE)
 # -----------------------------------
 st.sidebar.title("💪 FitPlan AI")
 page = st.sidebar.radio("Navigate", ["🏠 Fitness Form", "📊 BMI Result"])
 
+
 # -----------------------------------
 # PAGE 1: FITNESS FORM
 # -----------------------------------
@@ -59,7 +96,6 @@ if page == "🏠 Fitness Form":
 
     st.markdown("<div class='title'>FitPlan AI</div>", unsafe_allow_html=True)
     st.write("### Milestone 1: Fitness Profile & BMI Calculator")
-    st.write("")
 
     st.markdown("<div class='card'>", unsafe_allow_html=True)
 
@@ -87,7 +123,7 @@ if page == "🏠 Fitness Form":
     )
 
     equipment = st.multiselect(
-        "Available Equipment (Multiple Selection Allowed)",
+        "Available Equipment",
         ["Dumbbells", "Resistance Band", "Yoga Mat", "No Equipment"]
     )
 
@@ -95,35 +131,37 @@ if page == "🏠 Fitness Form":
 
     st.markdown("</div>", unsafe_allow_html=True)
 
+
+    # -------------------------
+    # MODIFY THIS PART
+    # -------------------------
     if submit:
 
-        # Validation
         if not name.strip():
             st.error("Name is required.")
+
         elif height_cm <= 0:
             st.error("Height must be greater than zero.")
+
         elif weight_kg <= 0:
             st.error("Weight must be greater than zero.")
+
         else:
-            height_m = height_cm / 100
-            bmi = weight_kg / (height_m ** 2)
-            bmi = round(bmi, 2)
-
-            if bmi < 18.5:
-                category = "Underweight"
-            elif 18.5 <= bmi < 24.9:
-                category = "Normal"
-            elif 25 <= bmi < 29.9:
-                category = "Overweight"
-            else:
-                category = "Obese"
-
-            # Save in session
+
+            bmi = calculate_bmi(weight_kg, height_cm)
+
+            category = get_category(bmi)
+
             st.session_state["name"] = name
             st.session_state["bmi"] = bmi
             st.session_state["category"] = category
+            st.session_state["goal"] = goal
+            st.session_state["level"] = fitness_level
+            st.session_state["equipment"] = equipment
+
+            st.success("BMI Calculated! Go to Result page.")
+
 
-            st.success("BMI Calculated! Go to 'BMI Result' page.")
 
 # -----------------------------------
 # PAGE 2: RESULT PAGE
@@ -131,11 +169,13 @@ if page == "🏠 Fitness Form":
 elif page == "📊 BMI Result":
 
     st.markdown("<div class='title'>Your BMI Report</div>", unsafe_allow_html=True)
-    st.write("")
 
     if "bmi" not in st.session_state:
+
         st.warning("Please fill the form first.")
+
     else:
+
         st.markdown(f"""
         <div class='result-card'>
             <h2>Hello, {st.session_state['name']} 👋</h2>
@@ -144,4 +184,30 @@ elif page == "📊 BMI Result":
         </div>
         """, unsafe_allow_html=True)
 
+
+        # -------------------------
+        # ADD THIS AI PART
+        # -------------------------
+
+        prompt = f"""
+        Create a 5 day workout plan.
+
+        BMI: {st.session_state['bmi']}
+        Category: {st.session_state['category']}
+        Goal: {st.session_state['goal']}
+        Fitness Level: {st.session_state['level']}
+        Equipment: {st.session_state['equipment']}
+        """
+
+        with st.spinner("Generating workout plan..."):
+
+            result = generator(
+                prompt,
+                max_new_tokens=300
+            )[0]["generated_text"]
+
+        st.subheader("🏋️ Your AI Workout Plan")
+
+        st.write(result)
+
         st.balloons()