Update app.py

app.py

@@ -2,16 +2,18 @@ import os
 import re
 import json
 import subprocess
-import cv2
-import numpy as np
+import tempfile
+import time
 import img2pdf
 import gradio as gr
 import google.generativeai as genai
 from markdown_pdf import MarkdownPdf, Section
 from pdf2image import convert_from_path
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
+import cv2
+import numpy as np
 
-# ---------- PROMPTS ----------
+# ---------- PROMPTS (preserved exactly) ----------
 PROMPTS = {
     "ALIGNMENT_PROMPT": {
         "role": "system",
@@ -88,9 +90,11 @@ Then show total clearly:
 }
 
 # -------------------- CONFIG --------------------
+# The Gemini API key must be set in the environment
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
-GRID_ROWS, GRID_COLS = 20, 14  # grid for imprint placement
+# Grid config for imprinting
+GRID_ROWS, GRID_COLS = 20, 14
 
 # ---------- HELPERS ----------
 def save_as_pdf(text, filename="output.pdf"):
@@ -100,14 +104,24 @@ def save_as_pdf(text, filename="output.pdf"):
     return filename
 
 def compress_pdf(input_path, output_path=None, max_size=20*1024*1024):
-    """Compress PDF only if larger than max_size (20MB default)."""
+    """
+    Compress PDF only if its size is larger than max_size (default 20MB).
+    Returns path to (possibly compressed) file.
+    """
     if output_path is None:
         base, ext = os.path.splitext(input_path)
         output_path = f"{base}_compressed{ext}"
 
-    if os.path.getsize(input_path) <= max_size:
+    try:
+        size = os.path.getsize(input_path)
+    except Exception:
+        return input_path
+
+    if size <= max_size:
+        print(f"ℹ️ Not compressing {input_path} ({size/1024/1024:.2f} MB <= {max_size/1024/1024} MB)")
         return input_path
 
+    print(f"🔎 Compressing {input_path} ({size/1024/1024:.2f} MB) -> {output_path}")
     try:
         gs_cmd = [
             "gs", "-sDEVICE=pdfwrite",
@@ -117,120 +131,341 @@ def compress_pdf(input_path, output_path=None, max_size=20*1024*1024):
             f"-sOutputFile={output_path}", input_path
         ]
         subprocess.run(gs_cmd, check=True)
-        if os.path.getsize(output_path) <= max_size:
+        new_size = os.path.getsize(output_path)
+        print(f"✅ Compression done. New size: {new_size/1024/1024:.2f} MB")
+        if new_size <= max_size:
             return output_path
         else:
+            print("⚠️ Compressed file still larger than threshold; returning original")
             return input_path
-    except Exception:
+    except Exception as e:
+        print("❌ Compression error:", e)
         return input_path
 
 def create_model():
     try:
+        print("⚡ Using gemini-2.5-pro model")
         return genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
     except Exception:
+        print("⚡ Falling back to gemini-2.5-flash model")
         return genai.GenerativeModel("gemini-2.5-flash", generation_config={"temperature": 0})
 
-# ---------- Extract marks per question ----------
+# ---------- Extract marks per question (parse grading Markdown) ----------
 def extract_marks_from_grading(grading_text):
+    """
+    Parse the grading markdown produced by the GRADING_PROMPT and extract marks per question.
+    Returns dict: {"grading": [{"question": "1.a", "marks_awarded": ["M1","A1"]}, ...]}
+    """
     grading_json = {"grading": []}
-    # Split by question sections
-    question_blocks = re.split(r"## Question\s+", grading_text)
-    for block in question_blocks[1:]:  # skip intro
-        # Extract question ID (like "1(a)" or "2.b")
-        q_match = re.match(r"([\d\.a-zA-Z\(\)]+)", block.strip())
-        if not q_match:
-            continue
-        q_id = q_match.group(1).strip()
-
-        # Find awarded marks in that block
+
+    # Split by question sections using "## Question" header
+    # We allow various header spacing, e.g. "## Question 1(a)" or "## Question 1(a)\n..."
+    question_blocks = re.split(r"##\s*Question\s+", grading_text)
+    for block in question_blocks[1:]:  # skip anything before first question header
+        # The first token up to newline is the question id line, e.g. "1(a)\n### Markscheme..."
+        first_line = block.strip().splitlines()[0].strip()
+        # Extract the question id - keep typical formats like 1, 1(a), 2.b, 3.d(ii)
+        q_id_match = re.match(r"([0-9]+(?:[a-zA-Z]|\([^\)]+\)|(?:\.[a-zA-Z0-9]+))*)", first_line)
+        if not q_id_match:
+            # fallback: try to extract tokens until first space
+            q_id = first_line.split()[0]
+        else:
+            q_id = q_id_match.group(1).strip()
+
+        # Now find all awarded marks in that block. Search the "Awarded" column entries like M1, A1, A0, R1 etc.
+        # We use a word-boundary regex to capture tokens.
         awarded = re.findall(r"\b(M\d+|A\d+|R\d+|M0|A0|R0)\b", block)
+        # Deduplicate preserving order
+        seen = set()
+        awarded_unique = []
+        for m in awarded:
+            if m not in seen:
+                awarded_unique.append(m)
+                seen.add(m)
 
         grading_json["grading"].append({
             "question": q_id,
-            "marks_awarded": awarded
+            "marks_awarded": awarded_unique
         })
 
     return grading_json
 
-# ---------- Imprinting Logic ----------
-def imprint_marks(pdf_path, grading_json, output_pdf):
+# ---------- Ask Gemini for mapping for one page image ----------
+def ask_gemini_for_mapping_for_page(model, image_path, grading_json, rows=GRID_ROWS, cols=GRID_COLS):
+    """
+    Sends a page image and grading JSON to Gemini asking for cell numbers for questions.
+    Returns a list like: [{"question": "1.a", "cell_number": 23}, ...]
+    """
+    prompt = f"""
+You are an exam marker. Your role is to identify where each question begins on the page.
+The page is divided into a {rows} x {cols} grid. Each cell has a RUNNING NUMBER label (1..{rows*cols}).
+For each question in the grading JSON, return the cell NUMBER where the FIRST STEP of that question begins.
+
+IMPORTANT RULES:
+- Do not place marks inside another question's answer area.
+- Prefer placing the marks in a BLANK cell immediately to the RIGHT of the answer step. If no blank cell is available to the right, then place in a blank cell to the LEFT.
+- Never place marks above or below the answer.
+- If a question starts on a previous page, you may omit it for this page.
+Return JSON only, like:
+[{{"question": "1.a", "cell_number": 15}}, ...]
+
+Grading JSON:
+{json.dumps(grading_json, indent=2)}
+"""
+    # Load image file
+    img = Image.open(image_path)
+    # Send both prompt and image to Gemini
+    response = model.generate_content([prompt, img])
+    raw_text = getattr(response, "text", None)
+    if not raw_text and getattr(response, "candidates", None):
+        raw_text = response.candidates[0].content.parts[0].text
+
+    print("🔎 Gemini mapping raw output (page):")
+    print(raw_text)
+
+    # Try to extract JSON from response
+    # Commonly model will return JSON; attempt to parse the first JSON array/list block
+    json_part = None
+    try:
+        # naive: find first '[' and last ']' and json.loads
+        start = raw_text.index('[')
+        end = raw_text.rindex(']') + 1
+        json_part = raw_text[start:end]
+        mapping = json.loads(json_part)
+        return mapping
+    except Exception as e:
+        print("⚠️ Failed to parse mapping JSON directly:", e)
+        # try to find 'json\n{...}\n' patterns
+        match = re.search(r'(\[.*\])', raw_text, re.DOTALL)
+        if match:
+            try:
+                mapping = json.loads(match.group(1))
+                return mapping
+            except Exception as e2:
+                print("⚠️ Second parse attempt failed:", e2)
+        # fallback empty list
+        return []
+
+# ---------- Imprinting Logic (uses mapping) ----------
+def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, rows=GRID_ROWS, cols=GRID_COLS):
+    """
+    Convert PDF to images, create grid-numbered images, ask Gemini for mapping per page,
+    and then annotate marks beside the mapped cells.
+    Returns path to final imprinted (and possibly compressed) PDF.
+    Prints imprint steps in console for each page/question.
+    """
     pages = convert_from_path(pdf_path, dpi=200)
-    annotated_pages = []
-
-    for idx, page in enumerate(pages):
-        img = np.array(page.convert("RGB"))
-        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-
-        y_offset = 100  # baseline vertical offset
-        for g in grading_json["grading"]:
-            marks_text = ",".join(g["marks_awarded"])
-            # Simple placement: stack vertically
-            cv2.putText(img, f"{g['question']}: {marks_text}",
-                        (50, y_offset),
-                        cv2.FONT_HERSHEY_SIMPLEX,
-                        1.2, (0, 0, 255), 3, cv2.LINE_AA)
-            y_offset += 50
-
-        annotated_path = f"annotated_{idx+1}.png"
-        cv2.imwrite(annotated_path, img)
-        annotated_pages.append(annotated_path)
-
+    annotated_page_paths = []
+    print(f"📄 Converted answer PDF to {len(pages)} page image(s) for imprinting.")
+
+    # create grid-numbered temporary images for sending to Gemini
+    temp_grid_images = []
+    for p_index, page in enumerate(pages):
+        img = page.convert("RGB")
+        w, h = img.size
+        cell_w, cell_h = w / cols, h / rows
+
+        draw = ImageDraw.Draw(img)
+        try:
+            num_font = ImageFont.truetype("arial.ttf", 16)
+        except Exception:
+            num_font = ImageFont.load_default()
+
+        cell_num = 1
+        # We only need numbers for clarity when sending to model (but we won't draw gridlines)
+        for r in range(rows):
+            for c in range(cols):
+                x = int(c * cell_w + cell_w / 2)
+                y = int(r * cell_h + cell_h / 2)
+                text = str(cell_num)
+                bbox = draw.textbbox((0, 0), text, font=num_font)
+                tw = bbox[2] - bbox[0]
+                th = bbox[3] - bbox[1]
+                draw.text((x - tw/2, y - th/2), text, fill="black", font=num_font)
+                cell_num += 1
+
+        temp_path = f"page_{p_index+1}_grid.png"
+        img.save(temp_path, "PNG")
+        temp_grid_images.append(temp_path)
+
+    # Now for each page, ask Gemini for mapping
+    for p_index, grid_img_path in enumerate(temp_grid_images):
+        print(f"\n🛰 Sending page {p_index+1} to Gemini for mapping...")
+        mapping = ask_gemini_for_mapping_for_page(model, grid_img_path, grading_json, rows, cols)
+        print(f"🔁 Parsed mapping for page {p_index+1}: {mapping}")
+
+        # Prepare a clean copy of the original page to annotate (no grid numbers)
+        page_img = pages[p_index].convert("RGB")
+        img_cv = np.array(page_img)
+        img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
+        h, w, _ = img_cv.shape
+        cell_w_px, cell_h_px = w / cols, h / rows
+
+        # We will maintain a set of occupied cells to prefer right/left placement heuristics
+        occupied = set()
+
+        # For each mapping entry, place the corresponding marks
+        for item in mapping:
+            qid = item.get("question")
+            cell_number = item.get("cell_number")
+            if qid is None or cell_number is None:
+                continue
+
+            # Find marks for this question from grading_json
+            marks_list = next((g["marks_awarded"] for g in grading_json.get("grading", []) if g["question"] == qid), [])
+            if not marks_list:
+                # possible the grading JSON uses slightly different formatting of q ids; try case-insensitive match
+                marks_list = next((g["marks_awarded"] for g in grading_json.get("grading", [])
+                                   if g["question"].lower() == qid.lower()), [])
+
+            marks_text = ",".join(marks_list) if marks_list else "?"
+
+            # Compute candidate cell coordinates
+            # Convert cell_number -> (row, col)
+            row = (cell_number - 1) // cols
+            col = (cell_number - 1) % cols
+
+            # Preference: place in cell to the right (col + 1), if within grid and not occupied.
+            placed = False
+            candidates = []
+            # Right cell
+            if col + 1 < cols:
+                candidates.append((row, col + 1))
+            # same cell (fallback)
+            candidates.append((row, col))
+            # left cell
+            if col - 1 >= 0:
+                candidates.append((row, col - 1))
+
+            chosen = None
+            for (r_c, c_c) in candidates:
+                cell_id = r_c * cols + c_c + 1
+                if cell_id not in occupied:
+                    chosen = (r_c, c_c)
+                    occupied.add(cell_id)
+                    break
+
+            if chosen is None:
+                # all occupied? just pick original cell
+                chosen = (row, col)
+
+            # Convert chosen cell to pixel coordinates (approx center-right)
+            r_c, c_c = chosen
+            x_c = int((c_c + 1) * cell_w_px - cell_w_px * 0.1)  # near right edge of cell
+            y_c = int((r_c + 0.5) * cell_h_px)
+
+            # Print the imprint step to console
+            print(f"Page {p_index+1} | Question {qid} -> mapped cell {cell_number} -> chosen cell ({r_c},{c_c})"
+                  f" -> pixel coords ({x_c},{y_c}) | marks: {marks_text}")
+
+            # Draw the text on the image (scale font according to cell size)
+            font_scale = max(0.6, min(1.6, cell_h_px / 60.0))
+            thickness = max(1, int(font_scale * 2))
+            # Use cv2.putText (BGR)
+            cv2.putText(img_cv, marks_text, (x_c, y_c), cv2.FONT_HERSHEY_SIMPLEX,
+                        font_scale, (0, 0, 255), thickness, cv2.LINE_AA)
+
+        # Save annotated page
+        annotated_path = f"annotated_page_{p_index+1}.png"
+        cv2.imwrite(annotated_path, img_cv)
+        annotated_page_paths.append(annotated_path)
+        print(f"🖊 Annotated page saved: {annotated_path}")
+
+    # Merge annotated pages into a PDF
     with open(output_pdf, "wb") as f:
-        f.write(img2pdf.convert(annotated_pages))
-
-    return compress_pdf(output_pdf)
-
-# ---------- PIPELINE ----------
-def align_and_grade(qp_file, ms_file, ans_file, imprint=False):
+        f.write(img2pdf.convert(annotated_page_paths))
+
+    print(f"📑 Imprinted PDF saved to: {output_pdf}")
+    # Compress output PDF only if > 20MB
+    compressed = compress_pdf(output_pdf)
+    if compressed != output_pdf:
+        print(f"📦 Imprinted PDF compressed: {compressed}")
+    return compressed
+
+# ---------- Main pipeline ----------
+def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
+    """
+    Runs: compress (if needed) -> upload files -> alignment -> grading -> extract marks ->
+          optional imprint (per-page mapping + annotation).
+    Returns: (alignment_json_text, grading_markdown_text, grading_pdf_path, imprinted_pdf_path or None)
+    """
     try:
-        # Compress only if >20MB
-        qp_file = compress_pdf(qp_file, "qp_compressed.pdf")
-        ms_file = compress_pdf(ms_file, "ms_compressed.pdf")
-        ans_file = compress_pdf(ans_file, "ans_compressed.pdf")
+        # Step 0: compress only if >20MB
+        qp_path = compress_pdf(qp_path)
+        ms_path = compress_pdf(ms_path)
+        ans_path = compress_pdf(ans_path)
 
-        qp_uploaded = genai.upload_file(path=qp_file, display_name="Question Paper")
-        ms_uploaded = genai.upload_file(path=ms_file, display_name="Markscheme")
-        ans_uploaded = genai.upload_file(path=ans_file, display_name="Answer Sheet")
+        # Upload files to Gemini
+        print("🔼 Uploading files to Gemini...")
+        qp_uploaded = genai.upload_file(path=qp_path, display_name="Question Paper")
+        ms_uploaded = genai.upload_file(path=ms_path, display_name="Markscheme")
+        ans_uploaded = genai.upload_file(path=ans_path, display_name="Answer Sheet")
 
         model = create_model()
 
-        # ---- Step 1: ALIGN (JSON only)
+        # Step 1: ALIGN (JSON only)
+        print("1️⃣ Sending ALIGNMENT_PROMPT to Gemini (alignment step)...")
         resp = model.generate_content([
             PROMPTS["ALIGNMENT_PROMPT"]["content"],
             qp_uploaded,
             ms_uploaded,
             ans_uploaded
         ])
+
         json_output = getattr(resp, "text", None)
-        if not json_output and resp.candidates:
+        if not json_output and getattr(resp, "candidates", None):
             json_output = resp.candidates[0].content.parts[0].text
 
-        # ---- Step 2: GRADING (Markdown)
+        # Ensure we have text
+        if not json_output:
+            raise RuntimeError("No alignment JSON returned from Gemini.")
+
+        print("✅ Alignment JSON received (truncated preview):")
+        print((json_output[:1000] + '...') if len(json_output) > 1000 else json_output)
+
+        # Step 2: GRADING (Markdown)
+        print("2️⃣ Sending GRADING_PROMPT to Gemini (grading step)...")
         response = model.generate_content([
             PROMPTS["GRADING_PROMPT"]["content"],
             json_output
         ])
-        grading = getattr(response, "text", None)
-        if not grading and response.candidates:
-            grading = response.candidates[0].content.parts[0].text
+        grading_text = getattr(response, "text", None)
+        if not grading_text and getattr(response, "candidates", None):
+            grading_text = response.candidates[0].content.parts[0].text
 
-        base_name = os.path.splitext(os.path.basename(ans_file))[0]
-        grading_pdf_path = save_as_pdf(grading, f"{base_name}_graded.pdf")
+        if not grading_text:
+            raise RuntimeError("No grading output returned from Gemini.")
 
-        # ---- Step 3 (Optional): Imprint marks on answer PDF ----
-        imprint_pdf_path = None
+        print("✅ Grading Markdown received (truncated preview):")
+        print((grading_text[:1000] + '...') if len(grading_text) > 1000 else grading_text)
+
+        # Save grading PDF
+        base_name = os.path.splitext(os.path.basename(ans_path))[0]
+        grading_pdf_path = save_as_pdf(grading_text, f"{base_name}_graded.pdf")
+        print(f"📄 Grading PDF saved: {grading_pdf_path}")
+
+        # Step 2.5: Extract marks per question from grading text
+        grading_json = extract_marks_from_grading(grading_text)
+        print("🔧 Extracted grading JSON (per-question marks):")
+        print(json.dumps(grading_json, indent=2))
+
+        imprinted_pdf_path = None
         if imprint:
-            grading_json = extract_marks_from_grading(grading)
-            imprint_pdf_path = imprint_marks(ans_file, grading_json, f"{base_name}_imprinted.pdf")
+            print("✍ Imprint option enabled. Starting imprinting process...")
+            # Convert answer PDF to grid pages, ask Gemini for mapping per page and annotate
+            imprinted_pdf_path = f"{base_name}_imprinted.pdf"
+            imprinted_pdf_path = imprint_marks_using_mapping(ans_path, grading_json, imprinted_pdf_path, model)
+            print(f"✅ Imprinting finished. Imprinted PDF at: {imprinted_pdf_path}")
 
-        return json_output, grading, grading_pdf_path, imprint_pdf_path
+        return json_output, grading_text, grading_pdf_path, imprinted_pdf_path
 
     except Exception as e:
+        print("❌ Pipeline error:", e)
         return f"❌ Error: {e}", None, None, None
 
 # ---------- GRADIO APP ----------
 with gr.Blocks(title="LeadIB AI Grading (Alignment + Auto-Grading + Imprint)") as demo:
-    gr.Markdown("## 📘 LeadIB AI Grading\nUpload **Question Paper**, **Markscheme**, and **Student Answer Sheet**.\nSystem aligns → grades → optionally imprints marks.")
+    gr.Markdown("## 📘 LeadIB AI Grading\nUpload **Question Paper**, **Markscheme**, and **Student Answer Sheet**.\nSystem aligns → grades → optionally imprints marks (per-question, per-page mapping).")
 
     with gr.Row():
         qp_file = gr.File(label="📄 Upload Question Paper (PDF)")
@@ -241,16 +476,29 @@ with gr.Blocks(title="LeadIB AI Grading (Alignment + Auto-Grading + Imprint)") a
     run_button = gr.Button("🚀 Run Alignment + Grading")
 
     with gr.Row():
-        json_output = gr.Textbox(label="📑 Step 1: Alignment (JSON)", lines=20)
-        grading_output = gr.Textbox(label="📝 Step 2: Grading (Markdown)", lines=20)
+        json_output_box = gr.Textbox(label="📑 Step 1: Alignment (JSON)", lines=20)
+        grading_output_box = gr.Textbox(label="📝 Step 2: Grading (Markdown)", lines=20)
+
+    grading_pdf_file = gr.File(label="📥 Download Grading PDF")
+    imprint_pdf_file = gr.File(label="📥 Download Imprinted PDF (Optional)")
+
+    def run_pipeline(qp_file_obj, ms_file_obj, ans_file_obj, imprint_flag):
+        # Gradio File objects have .name attribute when saved locally
+        qp_path = qp_file_obj.name
+        ms_path = ms_file_obj.name
+        ans_path = ans_file_obj.name
+
+        alignment_text, grading_text, grading_pdf_path, imprinted_pdf_path = align_and_grade_pipeline(
+            qp_path, ms_path, ans_path, imprint=imprint_flag
+        )
 
-    grading_pdf = gr.File(label="📥 Download Grading PDF")
-    imprint_pdf = gr.File(label="📥 Download Imprinted PDF (Optional)")
+        # For Gradio file outputs: return file paths (or None)
+        return alignment_text, grading_text, grading_pdf_path, imprinted_pdf_path
 
     run_button.click(
-        fn=align_and_grade,
+        fn=run_pipeline,
         inputs=[qp_file, ms_file, ans_file, imprint_toggle],
-        outputs=[json_output, grading_output, grading_pdf, imprint_pdf]
+        outputs=[json_output_box, grading_output_box, grading_pdf_file, imprint_pdf_file]
     )
 
 if __name__ == "__main__":