Update app.py

app.py

@@ -562,6 +562,32 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
                 })
         print("🔗 Graph bundles for grading:", graph_bundles)
 
+        # --- NEW: Separate and save graph pages as images ---
+        ms_images = convert_from_path(ms_path, dpi=200)
+        ans_images = convert_from_path(ans_path, dpi=200)
+        graph_image_paths = []  # To collect for grading prompt if needed
+        for bundle in graph_bundles:
+            q = bundle['question']
+            ms_page = bundle['ms_page']
+            as_page = bundle['as_page']
+            # Markscheme graph page
+            ms_img_path = f"ms_graph_Q{q}.png"
+            try:
+                ms_images[ms_page-1].save(ms_img_path)
+                print(f"🖼️ Separated Markscheme graph page for Q{q}: source=MS, page={ms_page}, output={ms_img_path}")
+                graph_image_paths.append((f"ms_graph_Q{q}", ms_img_path))
+            except Exception as e:
+                print(f"❌ Failed to save Markscheme graph page for Q{q}: {e}")
+            # Answer Sheet graph page
+            as_img_path = f"as_graph_Q{q}.png"
+            try:
+                ans_images[as_page-1].save(as_img_path)
+                print(f"🖼️ Separated Answer Sheet graph page for Q{q}: source=AS, page={as_page}, output={as_img_path}")
+                graph_image_paths.append((f"as_graph_Q{q}", as_img_path))
+            except Exception as e:
+                print(f"❌ Failed to save Answer Sheet graph page for Q{q}: {e}")
+        # --- END NEW ---
+
         # Step 4: Grading - send both transcripts to grading model, inject graph bundle info
         print("2) Preparing grading input and sending to Gemini for grading...")
         grading_input = (
@@ -581,7 +607,15 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
             grading_input += graph_note
 
         grading_prompt_system = PROMPTS["GRADING_PROMPT"]["content"]
-        grading_text = gemini_generate_content(model, grading_prompt_system + "\n\nPlease grade the following transcripts:\n" + grading_input)
+        # --- NEW: Attach graph images to grading prompt if graph bundles exist ---
+        if graph_bundles and graph_image_paths:
+            print(f"📎 Attaching {len(graph_image_paths)} graph images to grading prompt for visual grading.")
+            # If gemini_generate_content supports image_obj list, pass it here. Otherwise, just print for now.
+            # grading_text = gemini_generate_content(model, grading_prompt_system + "\n\nPlease grade the following transcripts:\n" + grading_input, image_obj=[p[1] for p in graph_image_paths])
+            grading_text = gemini_generate_content(model, grading_prompt_system + "\n\nPlease grade the following transcripts:\n" + grading_input)
+        else:
+            grading_text = gemini_generate_content(model, grading_prompt_system + "\n\nPlease grade the following transcripts:\n" + grading_input)
+        # --- END NEW ---
         print("🧾 Grading output received. Saving debug file: debug_grading.md")
         with open("debug_grading.md", "w", encoding="utf-8") as f:
             f.write(grading_text)