Update app.py

app.py

@@ -55,7 +55,7 @@ def parse_numbered_steps(text):
 if "steps" not in st.session_state:
     st.session_state.steps = []          # List[(int, str)]
 if "images" not in st.session_state:
-    st.session_state.images = {}         # Dict[int, PIL.Image]
+    st.session_state.images = {}         # Dict[int, PIL.Image] (illustrations per step)
 if "tools_list" not in st.session_state:
     st.session_state.tools_list = []     # List[str]
 if "current_step" not in st.session_state:
@@ -98,7 +98,7 @@ def detect_category_and_generate(uploaded_file, context_text):
     1) Use Gemini chat to identify the category (appliance, automotive, gardening, upcycling).
     2) Then use Gemini chat to get:
        - a "tools & materials" list
-       - numbered step-by-step instructions
+       - numbered step-by-step instructions (and optional illustrations)
     3) Parse everything into session_state.
     """
     try:
@@ -113,15 +113,13 @@ def detect_category_and_generate(uploaded_file, context_text):
             "home appliance repair, automotive maintenance, gardening & urban farming, or upcycling & sustainable crafts. "
             "Reply back with exactly one category name."
         )
-        # Create a chat instance that can handle both Text and Image
         cat_chat = client.chats.create(
             model=CATEGORY_MODEL,
             config=types.GenerateContentConfig(response_modalities=["Text", "Image"])
         )
-        # Send the category prompt plus the PIL.Image directly
         cat_resp = cat_chat.send_message([category_prompt, image])
 
-        # Extract the category name (assume it's the first text part)
+        # Extract the category name (combine all text parts)
         cat_parts = cat_resp.candidates[0].content.parts
         cat_text = ""
         for part in cat_parts:
@@ -143,6 +141,7 @@ def detect_category_and_generate(uploaded_file, context_text):
             "1. First step instructions (be specific and detailed)...\n"
             "2. Second step instructions...\n"
             "3. Continue with all necessary steps...\n\n"
+            "Additionally, for each step you may include a brief illustrative image. "
             "Keep each numbered step clear and actionable (2-3 sentences max). "
             "Include safety warnings where appropriate."
         )
@@ -152,20 +151,28 @@ def detect_category_and_generate(uploaded_file, context_text):
         )
         full_resp = gen_chat.send_message([detailed_prompt, image])
 
-        # ─── 3.3 PARSE out tools + numbered steps ─────────────────────────────────
-        # Combine all text parts from the response
+        # ─── 3.3 PARSE out tools, numbered steps, and images ────────────────────────
         gen_parts = full_resp.candidates[0].content.parts
-        response_text = ""
+
+        # First, gather all text into one string to parse tools/steps
+        combined_text = ""
         for part in gen_parts:
             if part.text is not None:
-                response_text += part.text + "\n"
-        response_text = response_text.strip()
+                combined_text += part.text + "\n"
+        combined_text = combined_text.strip()
+
+        # Now identify and extract any inline images. We'll assign them sequentially to steps.
+        inline_images = []
+        for part in gen_parts:
+            if part.inline_data is not None:
+                img = Image.open(BytesIO(part.inline_data.data))
+                inline_images.append(img)
 
-        # Split into tools and steps sections
+        # Split combined_text into tools and steps sections
         tools_section = ""
         steps_section = ""
-        if "TOOLS AND MATERIALS:" in response_text:
-            parts = response_text.split("TOOLS AND MATERIALS:")
+        if "TOOLS AND MATERIALS:" in combined_text:
+            parts = combined_text.split("TOOLS AND MATERIALS:")
             if len(parts) > 1:
                 remaining = parts[1]
                 if "STEPS:" in remaining:
@@ -174,13 +181,13 @@ def detect_category_and_generate(uploaded_file, context_text):
                     steps_section = steps_part.strip()
                 else:
                     tools_section = remaining.strip()
-        elif "STEPS:" in response_text:
-            parts = response_text.split("STEPS:")
+        elif "STEPS:" in combined_text:
+            parts = combined_text.split("STEPS:")
             if len(parts) > 1:
                 steps_section = parts[1].strip()
         else:
             # Fallback parsing
-            lines = response_text.split("\n")
+            lines = combined_text.split("\n")
             tools_lines = []
             steps_lines = []
             current_section = "unknown"
@@ -207,7 +214,14 @@ def detect_category_and_generate(uploaded_file, context_text):
         # Parse steps
         parsed_steps = parse_numbered_steps(steps_section)
 
-        # Store in session_state
+        # Assign inline_images to steps (one-to-one, up to number of steps)
+        st.session_state.images = {}
+        for idx, step in parsed_steps:
+            img_index = idx - 1  # zero-based
+            if img_index < len(inline_images):
+                st.session_state.images[idx] = inline_images[img_index]
+
+        # Store tools and steps in session_state
         st.session_state.tools_list = tools
         st.session_state.steps = parsed_steps
 
@@ -230,6 +244,7 @@ def detect_category_and_generate(uploaded_file, context_text):
         st.error(f"Error processing request: {str(e)}")
         st.error("Please check your API key and try again.")
 
+
 def render_sidebar_navigation():
     """
     A sidebar listing each step by number + a button to jump to that step.
@@ -242,13 +257,14 @@ def render_sidebar_navigation():
         st.sidebar.progress(completed / total_steps)
         st.sidebar.write(f"Progress: {completed}/{total_steps} steps")
 
-    for (idx, text) in st.session_state.steps:
+    for (idx, _) in st.session_state.steps:
         is_done = st.session_state.done_flags.get(idx, False)
         label = f"{'✓' if is_done else '·'} Step {idx}"
         if st.sidebar.button(label, key=f"nav_{idx}"):
             st.session_state.current_step = idx
             st.rerun()
 
+
 def render_tools_list():
     """Show the Tools/Materials list in an expander."""
     if st.session_state.tools_list:
@@ -256,11 +272,13 @@ def render_tools_list():
             for item in st.session_state.tools_list:
                 st.write(f"- {item}")
 
+
 def render_step(idx, text):
     """
     Render a single step:
     - Show step number
     - Show instruction text
+    - Possibly show an illustration if available
     - Timer if needed
     - Checkbox for "Done"
     - Text area for notes
@@ -271,6 +289,14 @@ def render_step(idx, text):
     st.markdown(f"### Step {idx} of {total}")
     st.write(text)
 
+    # If an illustration exists for this step, display it
+    if idx in st.session_state.images:
+        st.image(
+            st.session_state.images[idx],
+            caption=f"Illustration for step {idx}",
+            use_container_width=True
+        )
+
     # Timer functionality
     seconds_left = st.session_state.timers.get(idx, 0)
     if seconds_left > 0:
@@ -353,9 +379,10 @@ with st.expander("ℹ️ How it works", expanded=False):
         """
         1. **Upload a photo** of the item you want to fix or build (appliance, car part, plant, craft project).
         2. **Add context** (optional) - describe what’s wrong or what you want to achieve.
-        3. **Get AI guidance** - The AI will detect the category and provide step-by-step instructions.
+        3. **Get AI guidance** - The AI will detect the category and provide step-by-step instructions (with possible illustrations).
         4. **Follow the steps** - Each step includes:
            - Clear instructions
+           - An illustrative image (if provided)
            - Progress tracking with checkboxes
            - Timer functionality for waiting periods
            - Note-taking area
@@ -416,8 +443,7 @@ if st.session_state.prompt_sent:
         elif st.session_state.current_step < 1:
             st.session_state.current_step = 1
 
-        step_tuple = st.session_state.steps[st.session_state.current_step - 1]
-        step_num, step_text = step_tuple
+        step_num, step_text = st.session_state.steps[st.session_state.current_step - 1]
         render_step(step_num, step_text)
 
         # Overall progress at bottom