Update app.py

app.py

@@ -63,7 +63,7 @@ Answer 1.i:
 5. Apply FT where appropriate.  
 6. Use proper notation: M1A0, A1, etc.  
 7. Any lost mark: use red `<span style="color:red">M0</span>` and make Reason red.  
----  
+---
 ## Output Format
 Produce two sections per question/sub-question, following this structure:
 ## Question <id>
@@ -152,26 +152,24 @@ def gemini_generate_content(model, prompt_text, file_upload_obj=None, image_obj=
 # ---------------- PARSERS ----------------
 def extract_question_ids_from_qpms(text):
     """
-    Find all question ids in order *without* deduplication.
-    We will collect every match in sequence exactly as found.
+    Extract question IDs in the order they appear.
+    NOTE: do NOT deduplicate — keep multiple occurrences as they are in the transcript.
     """
     ids = []
-    # first try an explicit "Question: <id>" pattern
+    # first attempt: explicit "Question :" lines
     for m in re.finditer(r"(?im)^\s*Question\s*:\s*([0-9]+(?:[a-zA-Z0-9\.\(\)]+)*)\b", text):
         qid = m.group(1).strip()
         ids.append(qid)
-    # if none found by that pattern, use a looser leading numbering pattern
+    # fallback: lines starting with numbering like "1." or "2)" etc.
+    for m in re.finditer(r"(?m)^\s*([0-9]+(?:[a-zA-Z0-9\.\(\)]+)*)\s*[\.\):\-]\s", text):
+        qid = m.group(1).strip()
+        ids.append(qid)
+    # If nothing found, record "NA" once
     if not ids:
-        for m in re.finditer(r"(?m)^\s*([0-9]+(?:[a-zA-Z0-9\.\(\)]+)*)\s*[\.\):\-]\s", text):
-            qid = m.group(1).strip()
-            ids.append(qid)
-    return ids if ids else ["NA"]
+        ids = ["NA"]
+    return ids
 
 def build_as_prompt_with_expected_ids(expected_ids):
-    """
-    Build the AS transcription prompt; also useful to produce an ids_block string
-    that can be passed to the imprint mapping prompt.
-    """
     ids_block = "{\n" + "\n".join(expected_ids) + "\n}" if expected_ids else "{NA}"
     prompt = f"""You are a high-quality handwritten transcription assistant.
 INPUT: This PDF contains a student's handwritten answer sheet.
@@ -186,15 +184,22 @@ Question <id>
 AS:
 <transcribed answer or placeholder>
 """
-    return prompt, ids_block
+    return prompt
 
 def extract_marks_from_grading_exact(grading_text):
+    """
+    Extract grading marks in the order they appear and keep duplicates.
+    Output JSON with grading list preserving sequence (no deduplication).
+    """
     grading_json = {"grading": []}
+    # split by question blocks by heading "## Question "
     question_blocks = re.split(r"##\s*Question\s+", grading_text)
     for block in question_blocks[1:]:
+        # try to get the ID from the first line (robust)
         first_line = block.strip().splitlines()[0].strip() if block.strip().splitlines() else ""
         q_id_match = re.match(r"([0-9]+(?:[a-zA-Z]|\([^\)]+\)|(?:\.[a-zA-Z0-9]+))*)", first_line)
         q_id = q_id_match.group(1).strip() if q_id_match else first_line.split()[0] if first_line else ""
+        # find all mark tokens in order and preserve duplicates
         awarded = re.findall(r"\b(M\d+|A\d+|R\d+|M0|A0|R0)\b", block)
         grading_json["grading"].append({"question": q_id, "marks_awarded": awarded})
     return grading_json
@@ -202,39 +207,34 @@ def extract_marks_from_grading_exact(grading_text):
 # ---------------- IMPRINT ----------------
 def ask_gemini_for_mapping_for_page_v2(model, image_path, grading_json, question_scheme, ids_block, rows=GRID_ROWS, cols=GRID_COLS):
     """
-    Ask Gemini to map question IDs (from ids_block) to cell numbers on this page.
-    We pass the ids_block explicitly (same block used when transcribing student answers)
-    and instruct the model to return JSON only: a list of {"question":"<id>","cell_number":N}.
-
-    Also instruct the LLM about mislabelled subparts: e.g., if it sees 'ii)' above
-    'Q4.i' without a number, it may belong to Q3.ii.
+    Ask Gemini to map expected question IDs (ids_block) to grid cells on a single page image.
+    The prompt explicitly passes the expected IDs block and instructs the model to interpret
+    mislabelled steps (e.g., ii) above Q4 -> interpret as previous question's subpart).
     """
     prompt = f"""
-You are an exam marker. Identify where each question listed in the ids block begins on this page.
+You are an exam marker. Identify where each question begins on this page.
 The page has {rows}x{cols} grid (cells 1..{rows*cols}).
-QUESTION IDS (expected) you must look for:
+These are the QUESTIONS YOU MUST SPOT (expected IDs):
 {ids_block}
 
-Question scheme (authoritative transcription excerpt):
+Question scheme (authoritative full QP+MS text):
 {question_scheme}
 
-Grading JSON (marks awarded summary):
+Grading JSON:
 {json.dumps(grading_json, indent=2)}
 
-Instructions (IMPORTANT):
-- Only return questions from the provided IDs block above.
-- For each question you find on this page, return the single grid cell number where the first step of that question begins.
-- If you see a subpart like `ii)` with no leading question number directly above or below another labelled subpart, try to infer which question it belongs to (example: if you find `Q4.i` and above it you see `ii)` alone, it may be `Q3.ii` — if you infer like this explain your reasoning briefly in the JSON entry's optional "note" field).
-- Avoid placing marks inside another question's area; prefer an adjacent blank cell to the RIGHT if possible, else LEFT.
-- Return JSON only, exactly like:
-[
-  {"question":"1.a","cell_number":15},
-  {"question":"3.ii","cell_number":23, "note":"inferred from unlabeled ii above Q4.i"},
-  ...
-]
-- If no instances of an expected question appear on this page, return an empty list: [].
+Important instructions:
+- Only return cell numbers for the expected question IDs listed above.
+- If you detect a fragment like "ii)" above a later question heading (for example: you find "Q4" on the page and above it you see "ii)"), interpret that fragment as belonging to the previous question (e.g., "Q3.ii"). In other words: if a subpart appears spatially above a heading for the next question, reassign it to the previous question's appropriate subpart.
+- Return the earliest cell number where the student's first written step (the start of that answer) appears.
+- Prefer a blank cell immediately to the RIGHT of detected starting cell for placing marks; if not available, prefer LEFT.
+- Avoid placing marks inside another question's answer area where possible.
+- Only include questions that actually appear on this page.
+
+Return JSON only in the format:
+[{"question":"1.a","cell_number":15}, ...]
 """
-    # load image (PIL) so the model can see it if supported
+    # Attach image plus prompt to Gemini
     img = Image.open(image_path)
     response = model.generate_content([prompt, img])
     raw_text = getattr(response, "text", None)
@@ -242,39 +242,47 @@ Instructions (IMPORTANT):
         raw_text = response.candidates[0].content.parts[0].text
     if not raw_text:
         raw_text = str(response)
-    # try to extract a JSON array from the output
+    # attempt to parse JSON array from model output
     try:
         start = raw_text.index('[')
         end = raw_text.rindex(']') + 1
         return json.loads(raw_text[start:end])
     except Exception:
-        # if parsing fails, return an empty list for safety
-        return []
+        # Best-effort: try to extract lines like {"question":"1.a","cell_number":15}
+        try:
+            lines = re.findall(r'\{[^}]*\}', raw_text)
+            parsed = [json.loads(l) for l in lines]
+            return parsed
+        except Exception:
+            return []
 
-def imprint_marks_using_mapping_v2(pdf_path, grading_json, output_pdf, question_scheme, model, expected_ids, rows=GRID_ROWS, cols=GRID_COLS):
+def imprint_marks_using_mapping_v2(pdf_path, grading_json, output_pdf, question_scheme, model, ids_block, rows=GRID_ROWS, cols=GRID_COLS):
     """
-    Imprint marks onto the student answer PDF while preserving original page size.
-
-    - Reads page size from PDF (points).
-    - Converts pages to images with convert_from_path(..., size=(width_pt,height_pt)).
-    - Places grid and marks using cell widths/heights computed from the resulting image.
-    - Does NOT rescale images later.
-    - Writes final imprinted PDF using img2pdf with original page dimensions.
+    Imprint marks onto a PDF using mapping returned by Gemini.
+    Key changes:
+    - Use the PDF's original mediabox (width_pt, height_pt) and render pages at 72 DPI,
+      so that 1 pixel == 1 point and no scaling occurs.
+    - Create annotated images at native page size and recreate PDF using those exact dimensions.
+    - Print progress/log steps.
     """
+    print("[IMPRINT] Reading PDF and preparing page sizes...")
     reader = PdfReader(pdf_path)
-    # use the first page size as canonical for all pages (could be extended to per-page)
-    page0 = reader.pages[0]
-    width_pt = float(page0.mediabox.width)   # points
-    height_pt = float(page0.mediabox.height)
-
-    # convert PDF pages to images with exact size -> 1 image pixel ≈ 1 PDF point
-    # NOTE: pdf2image size expects a tuple of ints
-    pages = convert_from_path(pdf_path, size=(int(width_pt), int(height_pt)))
-
+    # Use first page mediabox as canonical (works if pages share same size). For multi-size PDFs,
+    # we will read each page size when processing that page.
+    pages_info = []
+    for p_index, p in enumerate(reader.pages):
+        width_pt = float(p.mediabox.width)
+        height_pt = float(p.mediabox.height)
+        pages_info.append({"index": p_index, "width_pt": width_pt, "height_pt": height_pt})
+
+    # Render pages at 72 DPI so pixel dimensions == points (1 pt = 1 px).
+    # This avoids any rescaling.
+    print("[IMPRINT] Converting PDF pages to images at 72 DPI (1 px == 1 point)...")
+    pages = convert_from_path(pdf_path, dpi=72)
     annotated_page_paths = []
     temp_grid_images = []
 
-    # Draw grid numbers on a copy (useful to send to LLM to ask mapping)
+    # Create grid overlays (for debugging/visual confirmation) and save images used for mapping
     for p_index, page_img in enumerate(pages):
         img = page_img.convert("RGB")
         draw = ImageDraw.Draw(img)
@@ -283,89 +291,123 @@ def imprint_marks_using_mapping_v2(pdf_path, grading_json, output_pdf, question_
         except:
             font = ImageFont.load_default()
 
-        cell_w = img.width / cols
-        cell_h = img.height / rows
+        cols_local = cols
+        rows_local = rows
+        cell_w = img.width / cols_local
+        cell_h = img.height / rows_local
         cell_num = 1
-        for r in range(rows):
-            for c in range(cols):
-                # center of cell
+        for r in range(rows_local):
+            for c in range(cols_local):
                 x = int(c * cell_w + cell_w / 2)
                 y = int(r * cell_h + cell_h / 2)
                 bbox = draw.textbbox((0,0), str(cell_num), font=font)
                 draw.text((x - (bbox[2]-bbox[0])/2, y - (bbox[3]-bbox[1])/2), str(cell_num), fill="black", font=font)
-                cell_num +=1
+                cell_num += 1
         grid_path = f"page_{p_index+1}_grid.png"
         img.save(grid_path, "PNG")
         temp_grid_images.append(grid_path)
+        print(f"[IMPRINT] Grid image saved: {grid_path} (pixels: {img.width}x{img.height})")
 
-    # Build ids_block from expected_ids
-    ids_block = "{\n" + "\n".join(expected_ids) + "\n}" if expected_ids else "{NA}"
-
-    # Ask model to map each page (parallel)
+    # Ask Gemini (concurrently) to map question starts to cells
+    print("[IMPRINT] Sending grid images to Gemini to obtain cell mappings...")
     mappings_per_page = {}
-    with ThreadPoolExecutor(max_workers=min(8,len(temp_grid_images))) as ex:
+    with ThreadPoolExecutor(max_workers=min(8, len(temp_grid_images))) as ex:
         futures = {ex.submit(ask_gemini_for_mapping_for_page_v2, model, img_path, grading_json, question_scheme, ids_block, rows, cols): idx
-                   for idx,img_path in enumerate(temp_grid_images)}
+                   for idx, img_path in enumerate(temp_grid_images)}
         for fut in as_completed(futures):
             idx = futures[fut]
             try:
-                mappings_per_page[idx] = fut.result()
-            except:
+                mapping_result = fut.result()
+                mappings_per_page[idx] = mapping_result
+                print(f"[IMPRINT] Mapping received for page {idx+1}: {mapping_result}")
+            except Exception as e:
                 mappings_per_page[idx] = []
+                print(f"[IMPRINT] Mapping failed for page {idx+1}: {e}")
 
-    # Annotate original pages (no rescaling)
+    # Now annotate pages with marks text using the mapping results
+    print("[IMPRINT] Annotating pages with marks...")
     for p_index, page_img in enumerate(pages):
         img_cv = np.array(page_img.convert("RGB"))
         img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
         h, w, _ = img_cv.shape
-        cell_w_px, cell_h_px = w/cols, h/rows
+        cell_w_px, cell_h_px = w / cols, h / rows
         mapping = mappings_per_page.get(p_index, [])
         occupied = set()
         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 that question in grading_json (exact string match)
-            marks_list = next((g["marks_awarded"] for g in grading_json.get("grading", []) if g["question"]==qid), [])
+            if qid is None or cell_number is None:
+                continue
+            marks_list = next((g["marks_awarded"] for g in grading_json.get("grading", []) if g["question"] == qid), [])
             marks_text = ",".join(marks_list) if marks_list else "?"
-            # map cell_number -> row/col
-            row = (cell_number-1)//cols
-            col = (cell_number-1)%cols
-            # choose preference: right, same, left
+            # compute canonical cell row/col
+            row = (cell_number - 1) // cols
+            col = (cell_number - 1) % cols
+            # candidate placements (prefer right, then same, then left)
             candidates = []
-            if col+1<cols: candidates.append((row,col+1))
-            candidates.append((row,col))
-            if col-1>=0: candidates.append((row,col-1))
-            chosen = next(((r,c) for r,c in candidates if (r*cols+c+1) not in occupied), (row,col))
-            occupied.add(chosen[0]*cols+chosen[1]+1)
-            x_c = int((chosen[1]+0.5)*cell_w_px)
-            y_c = int((chosen[0]+0.5)*cell_h_px)
-            # draw mark text directly onto image (OpenCV uses BGR)
-            font_scale = max(0.6,min(1.6,cell_h_px/60))
-            thickness = max(1,int(font_scale*2))
-            # ensure text doesn't go out of bounds; shift up a bit
-            text_size = cv2.getTextSize(marks_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)[0]
-            x_draw = max(0, min(w - text_size[0], x_c - text_size[0]//2))
-            y_draw = max(text_size[1], min(h - 1, y_c + text_size[1]//2))
-            cv2.putText(img_cv, marks_text, (x_draw,y_draw), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0,0,255), thickness)
-
-            # optional: if mapping includes "note", draw a small 'i' icon nearby
-            if item.get("note"):
-                note_text = "i"
-                ns = cv2.getTextSize(note_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale*0.8, max(1,int(thickness/2)))[0]
-                nx = max(0, x_draw + text_size[0] + 4)
-                ny = max(ns[1], y_draw)
-                cv2.putText(img_cv, note_text, (nx, ny), cv2.FONT_HERSHEY_SIMPLEX, font_scale*0.8, (0,0,0), max(1,int(thickness/2)))
-
+            if col + 1 < cols: candidates.append((row, col + 1))
+            candidates.append((row, col))
+            if col - 1 >= 0: candidates.append((row, col - 1))
+            chosen = next(((r, c) for r, c in candidates if (r * cols + c + 1) not in occupied), (row, col))
+            occupied.add(chosen[0] * cols + chosen[1] + 1)
+            x_c = int((chosen[1] + 0.5) * cell_w_px)
+            y_c = int((chosen[0] + 0.5) * cell_h_px)
+            font_scale = max(0.6, min(1.6, cell_h_px / 60))
+            thickness = max(1, int(font_scale * 2))
+            cv2.putText(img_cv, marks_text, (x_c, y_c), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 255), thickness)
+            print(f"[IMPRINT] Placed marks '{marks_text}' for '{qid}' at page {p_index+1} cell {cell_number} -> pixel ({x_c},{y_c})")
         annotated_path = f"annotated_page_{p_index+1}.png"
         cv2.imwrite(annotated_path, img_cv)
         annotated_page_paths.append(annotated_path)
-
-    # Recreate PDF using img2pdf with the original page dimensions (width_pt,height_pt)
-    with open(output_pdf,"wb") as f:
-        f.write(img2pdf.convert(annotated_page_paths, layout_fun=img2pdf.get_layout_fun((width_pt,height_pt))))
-
-    return compress_pdf(output_pdf)
+        print(f"[IMPRINT] Annotated page saved: {annotated_path}")
+
+    # Recreate PDF using img2pdf with the original page dimensions (points).
+    # Since we rendered at 72 DPI, pixels == points and using layout_fun with (width_pt, height_pt) will preserve size.
+    print("[IMPRINT] Recreating PDF from annotated pages with original page sizes...")
+    layout_sizes = []
+    for p_info in pages_info:
+        layout_sizes.append((p_info["width_pt"], p_info["height_pt"]))
+    # If every page has same mediabox, img2pdf.get_layout_fun can be given that size; otherwise fallback to a per-image function.
+    try:
+        # We will use the mediabox of the first page for layout function if single size, else create per-image layout
+        unique_sizes = { (p["width_pt"], p["height_pt"]) for p in pages_info }
+        if len(unique_sizes) == 1:
+            w_pt, h_pt = pages_info[0]["width_pt"], pages_info[0]["height_pt"]
+            with open(output_pdf, "wb") as f:
+                f.write(img2pdf.convert(annotated_page_paths, layout_fun=img2pdf.get_layout_fun((w_pt, h_pt))))
+        else:
+            # per-page layout: build a custom layout function for each image based on index
+            # img2pdf allows layout_fun that takes (img_width_px, img_height_px, px_density) but easier approach:
+            # create PDF by converting each annotated PNG individually to single-page PDF with proper size and then merge
+            per_page_pdfs = []
+            for idx, ann_path in enumerate(annotated_page_paths):
+                w_pt = pages_info[idx]["width_pt"]
+                h_pt = pages_info[idx]["height_pt"]
+                single_pdf = f"single_{idx+1}.pdf"
+                with open(single_pdf, "wb") as f:
+                    f.write(img2pdf.convert(ann_path, layout_fun=img2pdf.get_layout_fun((w_pt, h_pt))))
+                per_page_pdfs.append(single_pdf)
+            # merge them
+            merge_pdfs(per_page_pdfs, output_pdf)
+            # cleanup single_page temp pdfs
+            for p in per_page_pdfs:
+                try:
+                    os.remove(p)
+                except:
+                    pass
+    except Exception as e:
+        print(f"[IMPRINT] Failed to create imprinted PDF with original sizes: {e}")
+        # fallback: create naive pdf (may be resized)
+        with open(output_pdf, "wb") as f:
+            f.write(img2pdf.convert(annotated_page_paths))
+    print(f"[IMPRINT] Imprinted PDF created: {output_pdf}")
+
+    # Optionally compress result
+    compressed = compress_pdf(output_pdf)
+    if compressed != output_pdf:
+        print(f"[IMPRINT] Compressed imprinted PDF saved as: {compressed}")
+        return compressed
+    return output_pdf
 
 # ---------------- PIPELINE ----------------
 def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
@@ -373,7 +415,7 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
     ms_path = compress_pdf(ms_path)
     ans_path = compress_pdf(ans_path)
 
-    merged_qpms_path = os.path.splitext(qp_path)[0]+"_merged_qp_ms.pdf"
+    merged_qpms_path = os.path.splitext(qp_path)[0] + "_merged_qp_ms.pdf"
     merge_pdfs([qp_path, ms_path], merged_qpms_path)
 
     merged_uploaded = genai.upload_file(path=merged_qpms_path, display_name="QP+MS (merged)")
@@ -382,33 +424,42 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
     model = create_model()
 
     qpms_prompt = PROMPTS["QP_MS_TRANSCRIPTION"]["content"]
+    print("[STEP] Requesting QP+MS transcription from Gemini...")
     qpms_text = gemini_generate_content(model, qpms_prompt, file_upload_obj=merged_uploaded)
+    print("[STEP] QP+MS transcription received.")
 
-    # extract question ids (no deduplication)
     extracted_ids = extract_question_ids_from_qpms(qpms_text)
+    print(f"[STEP] Extracted question IDs (in order, duplicates preserved): {extracted_ids}")
 
-    # build AS prompt (and get ids_block)
-    as_prompt, ids_block = build_as_prompt_with_expected_ids(extracted_ids)
+    as_prompt = build_as_prompt_with_expected_ids(extracted_ids)
+    print("[STEP] Requesting AS transcription from Gemini (using expected IDs block)...")
     as_text = gemini_generate_content(model, as_prompt, file_upload_obj=ans_uploaded)
+    print("[STEP] AS transcription received.")
 
     grading_input = (
-        "=== QP+MS TRANSCRIPT BEGIN ===\n"+qpms_text+
-        "\n=== QP+MS TRANSCRIPT END ===\n\n"+
-        "=== ANSWER SHEET TRANSCRIPT BEGIN ===\n"+as_text+
+        "=== QP+MS TRANSCRIPT BEGIN ===\n" + qpms_text +
+        "\n=== QP+MS TRANSCRIPT END ===\n\n" +
+        "=== ANSWER SHEET TRANSCRIPT BEGIN ===\n" + as_text +
         "\n=== ANSWER SHEET TRANSCRIPT END ===\n"
     )
     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)
+    print("[STEP] Sending transcripts to Gemini for grading...")
+    grading_text = gemini_generate_content(model, grading_prompt_system + "\n\nPlease grade the following transcripts:\n" + grading_input)
+    print("[STEP] Grading received from Gemini.")
 
-    grading_pdf_path = save_as_pdf(grading_text, os.path.splitext(os.path.basename(ans_path))[0]+"_graded.pdf")
+    grading_pdf_path = save_as_pdf(grading_text, os.path.splitext(os.path.basename(ans_path))[0] + "_graded.pdf")
     grading_json = extract_marks_from_grading_exact(grading_text)
+    print(f"[STEP] Extracted grading JSON (duplicates preserved): {json.dumps(grading_json, indent=2)}")
 
     imprinted_pdf_path = None
     if imprint:
         question_scheme = qpms_text
-        imprinted_pdf_path = os.path.splitext(os.path.basename(ans_path))[0]+"_imprinted.pdf"
-        # Pass expected_ids (extracted_ids) to imprint function so it can build the ids_block and ask the model
-        imprinted_pdf_path = imprint_marks_using_mapping_v2(ans_path, grading_json, imprinted_pdf_path, question_scheme, model, expected_ids=extracted_ids)
+        imprinted_pdf_path = os.path.splitext(os.path.basename(ans_path))[0] + "_imprinted.pdf"
+        # Build ids_block to pass to ask_gemini_for_mapping_for_page_v2
+        ids_block = "{\n" + "\n".join(extracted_ids) + "\n}"
+        print("[IMPRINT] Starting imprinting with ids_block and question scheme...")
+        imprinted_pdf_path = imprint_marks_using_mapping_v2(ans_path, grading_json, imprinted_pdf_path, question_scheme, model, ids_block)
+        print(f"[IMPRINT] Completed imprinting. File: {imprinted_pdf_path}")
 
     return qpms_text, as_text, grading_text, grading_pdf_path, imprinted_pdf_path