Home update

Home.py

@@ -1,52 +1,102 @@
 import streamlit as st
 from PIL import Image
-from model.classifier import get_model, predict
+from model.classifier import get_model, predict, get_model_by_name
 from model.search_script import search_for_recipes
 import streamlit.components.v1 as components
-import time
 import base64
-
+import config as config
 from utils.layout import render_layout
 
+MODEL_PATH_MAP = {
+    "Onion": config.MODEL_PATH_ONION,
+    "Pear": config.MODEL_PATH_PEAR,
+    "Strawberry": config.MODEL_PATH_STRAWBERRY,
+    "Tomato": config.MODEL_PATH_TOMATO
+}
+
+VARIATION_CLASS_MAP = {
+    "Onion": ['halved', 'sliced', 'whole'],
+    "Strawberry": ['Hulled', 'sliced', 'whole'],
+    "Tomato": ['diced', 'vines', 'whole'],
+    "Pear": ['halved', 'sliced', 'whole']
+}
+
 @st.cache_resource
 def load_model():
     return get_model()
 
+@st.cache_resource
+def load_model_variation(product_name):
+    model_path = MODEL_PATH_MAP[product_name]
+    num_classes = len(VARIATION_CLASS_MAP[product_name])
+    return get_model_by_name(model_path, num_classes=num_classes)
+
 def classification_and_recommendation_page():
-    st.markdown("## 🖼️ Task A: Image Classification + 🍽️ Recipe Recommendation")
+    st.markdown("## 🍽️ Recipe Recommendation System")
     st.markdown("""
-    <div class="about-box">
+    <div style='
+        background-color: #f9f9f9;
+        border-left: 6px solid #4CAF50;
+        padding: 16px;
+        border-radius: 10px;
+        font-size: 15px;
+        line-height: 1.6;
+    '>
+    <b>📚 Recipe Recommendation Guide</b><br><br>
+
     Upload one or more food images. This module classifies each image into 
-    <b>Onion, Pear, Strawberry, or Tomato</b> using EfficientNet-B0, and then recommends recipes
-    based on the combined classification results.
-    </div>
+    <b>Onion, Pear, Strawberry, or Tomato</b> using <b>EfficientNet-B0</b>, and recommends recipes
+    based on the combined classification results.<br><br>
+
+    <b>Steps:</b><br>
+    1️⃣ Upload images (single or multiple) of produce, or directly add tags for recipe search.<br>
+    2️⃣ Once uploaded, the corresponding produce tag will be automatically added to the search.<br>
+    3️⃣ Use the sliders to choose the number of results and minimum recipe rating.<br>
+    4️⃣ Click <b>"Search Recipe"</b> to view personalized recommendations.
+    </div></br>                       
     """, unsafe_allow_html=True)
 
+
+
     model = load_model()
 
-    # --- Upload and classify ---
     uploaded_files = st.file_uploader("📤 Upload images (JPG/PNG)", type=["jpg", "jpeg", "png"], accept_multiple_files=True)
 
     if "uploaded_images" not in st.session_state:
         st.session_state.uploaded_images = []
     if "image_tags" not in st.session_state:
         st.session_state.image_tags = {}
+    if "image_variations" not in st.session_state:
+        st.session_state.image_variations = {}
 
     if uploaded_files:
         for img_file in uploaded_files:
             if img_file.name not in [img.name for img in st.session_state.uploaded_images]:
                 img = Image.open(img_file).convert("RGB")
-                label, _ = predict(img, model)
+                label, main_class_prob = predict(img, model)
+
+                variation = None
+                if label in VARIATION_CLASS_MAP:
+                    variation_model = load_model_variation(label)
+                    class_labels = VARIATION_CLASS_MAP[label]
+                    variation_label, var_conf = predict(img, variation_model, class_labels=class_labels)
+                    variation = f"{variation_label} ({var_conf*main_class_prob* 100:.1f}%)"
+
                 st.session_state.uploaded_images.append(img_file)
                 st.session_state.image_tags[img_file.name] = label
+                st.session_state.image_variations[img_file.name] = variation
 
-    # --- Show grid of classified images ---
+    current_file_names = [f.name for f in uploaded_files] if uploaded_files else []
+    st.session_state.uploaded_images = [f for f in st.session_state.uploaded_images if f.name in current_file_names]
+    st.session_state.image_tags = {k: v for k, v in st.session_state.image_tags.items() if k in current_file_names}
+    st.session_state.image_variations = {k: v for k, v in st.session_state.image_variations.items() if k in current_file_names}
+    
     if st.session_state.uploaded_images:
         html = """
         <style>
         .image-grid { display: flex; flex-wrap: wrap; gap: 12px; margin-top: 10px; }
         .image-card {
-            width: 140px; height: 180px;
+            width: 140px; height: 200px;
             border: 1px solid #ccc; border-radius: 10px;
             overflow: hidden; text-align: center;
             font-size: 13px; position: relative;
@@ -56,31 +106,28 @@ def classification_and_recommendation_page():
             max-width: 100%; max-height: 110px;
             object-fit: contain; margin-top: 5px;
         }
-        .remove-btn {
-            position: absolute; top: 2px; right: 6px;
-            color: #d33; background: #fff;
-            border: none; cursor: pointer; font-size: 16px;
-        }
         </style>
         <div class="image-grid">
         """
 
         for img in st.session_state.uploaded_images:
             label = st.session_state.image_tags.get(img.name, "unknown")
+            variation = st.session_state.image_variations.get(img.name, "")
+            combined_label = f"{label.upper()} </br> {variation}" if variation else label.upper()
             img_b64 = base64.b64encode(img.getvalue()).decode()
+
             html += f"""
             <div class="image-card">
                 <img src="data:image/png;base64,{img_b64}" />
-                <div><b>{label.upper()}</b></div>
+                <div style="margin-top: 5px; font-weight: bold; font-size: 13px;">{combined_label}</div>
                 <div style="color:gray; font-size:11px;">{img.name}</div>
             </div>
             """
 
         html += "</div>"
         grid_rows = ((len(st.session_state.uploaded_images) - 1) // 5 + 1)
-        components.html(html, height=200 * grid_rows + 20, scrolling=True)
+        components.html(html, height=200 * grid_rows + 40, scrolling=True)
 
-    # --- Recipe Search ---
     st.markdown("---")
     st.markdown("## 🔍 Recipe Recommendation")
 

config.py

@@ -1,4 +1,4 @@
-CLASS_LABELS = ['onion', 'pear', 'strawberry', 'tomato']
+CLASS_LABELS = ['Onion', 'Pear', 'Strawberry', 'Tomato']
 
 MODEL_PATH = "assets/modelWeights/best_model_v1.pth"
 MODEL_PATH_ONION = "assets/modelWeights/best_model_onion_v1.pth"