Update app.py

app.py

@@ -338,71 +338,76 @@ def extract_marks_from_grading(grading_text):
     return grading_json 
 
 # ---------------- MAPPING/IMPRINT HELPERS ---------------- 
-def ask_gemini_for_mapping_for_page(model, image_path, grading_json, expected_ids=None, rows=GRID_ROWS, cols=GRID_COLS): 
-    """ 
-    Send a single page image along with the grading_json and expected_ids; LLM should return JSON mapping. 
-    """ 
-    ids_block = "{NA}" 
-    if expected_ids: 
-        ids_block = "{\n" + "\n".join(expected_ids) + "\n}" 
-    prompt = f"""
-You are an exam marker. Your task is to locate a blank cell adjacent to the answer step and place the marks there:
-Primary preference: Use the blank cell immediately to the right of the answer step.
-Fallback: If no blank cell is available on the right, use the blank cell immediately to the left..
-The page is divided into a {rows} x {cols} grid. Each cell has a RUNNING NUMBER label (1..{rows*cols}).
+def ask_gemini_for_mapping_batch(model, image_paths, grading_json, expected_ids=None, rows=GRID_ROWS, cols=GRID_COLS):
+    """
+    Send multiple page images together to Gemini for batch mapping processing.
+    More efficient than sending one by one.
+    """
+    ids_block = "{NA}"
+    if expected_ids:
+        ids_block = "{\n" + "\n".join(expected_ids) + "\n}"
+    
+    prompt = f"""You are an exam marker. Your role is to identify where each question begins on each page.
+The pages are divided into a {rows} x {cols} grid. Each cell has a RUNNING NUMBER label.
 For each question in the grading JSON, return the cell NUMBER where the FIRST STEP of that question begins.
-IMPORTANT: For your help i have provided u questions that u can expect in the image  :
-{ids_block}
-If you see a sub-question (e.g., ii) above a main question (e.g., Q4), infer it belongs to the previous question (e.g., Q3.ii).
+
+⚠ IMPORTANT RULES:
 - Do not place marks inside another question's answer area.
-- Each question should have unique cell number 
-- If a question serial number is visible in the answer image, you must mandatorily identify the corresponding question using the grading JSON.
+- 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.
-- Only if there is no serial number u may omit to select cell number for mark placement 
+- Each question should have unique cell number
+- If a question serial number is visible in the answer image, you must mandatorily identify the corresponding question using the grading JSON.
+
+IMPORTANT: For your help i have provided u questions that u can expect in the images:
+{ids_block}
+
 Return JSON only, like:
-[{{"question": "1.a", "cell_number": 15}}, ...]
+[{{"page": 1, "question": "1(a)", "cell_number": 15}}, ...]
+
 Grading JSON:
-{json.dumps(grading_json, indent=2)}
-""" 
-    print(f"📡 Sending mapping request for image {image_path} to Gemini...") 
-    img = Image.open(image_path) 
-    response = model.generate_content([prompt, img])
-    print("💬 Gemini response:", response)
-    raw_text = getattr(response, "text", None) 
-    if not raw_text and getattr(response, "candidates", None): 
-        raw_text = response.candidates[0].content.parts[0].text 
-    if not raw_text: 
-        raw_text = str(response) 
-    print("📥 Mapping response (chars):", len(raw_text)) 
-    try: 
-        start = raw_text.index('[') 
-        end = raw_text.rindex(']') + 1 
-        json_part = raw_text[start:end] 
-        mapping = json.loads(json_part) 
-        print("✅ Parsed mapping JSON for", image_path, "| entries:", len(mapping)) 
-        return mapping 
-    except Exception: 
-        match = re.search(r'(\[.*\])', raw_text, re.DOTALL) 
-        if match: 
-            try: 
-                mapping = json.loads(match.group(1)) 
-                print("✅ Parsed mapping JSON (alt) for", image_path, "| entries:", len(mapping)) 
-                return mapping 
-            except Exception: 
-                pass 
-        print("⚠️ Failed to parse mapping JSON for", image_path) 
+{json.dumps(grading_json, indent=2)}"""
+
+    # Load all images
+    images = [Image.open(p) for p in image_paths]
+    
+    print(f"📡 Sending batch mapping request for {len(image_paths)} pages to Gemini...")
+    response = model.generate_content([prompt, *images])
+    
+    raw_text = getattr(response, "text", None)
+    if not raw_text and getattr(response, "candidates", None):
+        raw_text = response.candidates[0].content.parts[0].text
+    if not raw_text:
+        raw_text = str(response)
+    
+    print("📥 Batch mapping response (chars):", len(raw_text))
+    print("🔎 Gemini raw batch output:")
+    print(raw_text)
+    
+    # Try to extract JSON from response
+    try:
+        match = re.search(r'(\[.*\])', raw_text, re.DOTALL)
+        if match:
+            mapping = json.loads(match.group(1))
+            print(f"✅ Parsed Gemini batch mapping for {len(image_paths)} pages")
+            return mapping
+        else:
+            print("❌ Failed to find JSON array in response")
+            return []
+    except Exception as e:
+        print(f"❌ Failed to parse Gemini JSON mapping: {e}")
         return [] 
 
 def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, expected_ids=None, rows=GRID_ROWS, cols=GRID_COLS): 
     """ 
-    Convert PDF to images, create grid-numbered images for sending to Gemini, 
-    send all page images in parallel to Gemini for mapping, then annotate and produce imprinted PDF. 
+    Convert PDF to images, create grid-numbered images for batch sending to Gemini, 
+    then annotate and produce imprinted PDF using batch processing for better efficiency.
     """ 
     print("📄 Converting answer PDF to images for imprinting...") 
     pages = convert_from_path(pdf_path, dpi=200) 
     annotated_page_paths = [] 
     temp_grid_images = [] 
 
+    # Create grid images for Gemini
     for p_index, page in enumerate(pages): 
         img = page.convert("RGB") 
         w, h = img.size 
@@ -410,7 +415,7 @@ def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, expec
 
         draw = ImageDraw.Draw(img) 
         try: 
-            num_font = ImageFont.truetype("arial.ttf", 16) 
+            num_font = ImageFont.truetype("arial.ttf", 20) 
         except Exception: 
             num_font = ImageFont.load_default() 
 
@@ -431,34 +436,31 @@ def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, expec
         temp_grid_images.append(temp_path) 
         print("🛰 Created grid image:", temp_path) 
 
-    # Send all grid images in parallel to Gemini to get mappings 
-    print("📡 Sending all page images to Gemini in parallel for mapping...") 
-    mappings_per_page = {} 
-    model_local = model 
-    with ThreadPoolExecutor(max_workers=min(8, len(temp_grid_images))) as ex: 
-        futures = {ex.submit(ask_gemini_for_mapping_for_page, model_local, img_path, grading_json, expected_ids, rows, cols): idx 
-                   for idx, img_path in enumerate(temp_grid_images)} 
-        for fut in as_completed(futures): 
-            idx = futures[fut] 
-            try: 
-                mapping = fut.result() 
-            except Exception as e: 
-                print("⚠️ Mapping request failed for page", idx, e) 
-                mapping = [] 
-            mappings_per_page[idx] = mapping 
+    # Send pages in batches to Gemini for mapping
+    print("📡 Sending page images to Gemini in batches for mapping...")
+    batch_size = 10  # Process 10 pages at a time
+    all_mappings = []
+    
+    for start in range(0, len(temp_grid_images), batch_size):
+        batch_paths = temp_grid_images[start:start+batch_size]
+        batch_mapping = ask_gemini_for_mapping_batch(model, batch_paths, grading_json, expected_ids, rows, cols)
+        all_mappings.extend(batch_mapping)
+        print(f"✅ Processed batch {start//batch_size + 1}: pages {start+1}-{start+len(batch_paths)}")
 
     # Annotate original pages according to returned mappings 
     print("🖊 Annotating pages with marks...") 
     for p_index, page in enumerate(pages): 
+        page_num = p_index + 1
         page_img = page.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 
 
-        mapping = mappings_per_page.get(p_index, []) 
-        occupied = set() 
-        for item in mapping: 
+        # Filter mappings for this page
+        page_mappings = [m for m in all_mappings if m.get("page") == page_num]
+        
+        for item in page_mappings: 
             qid = item.get("question") 
             cell_number = item.get("cell_number") 
             if qid is None or cell_number is None: 
@@ -474,37 +476,24 @@ def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, expec
             row = (cell_number - 1) // cols 
             col = (cell_number - 1) % cols 
 
-            candidates = [] 
-            if col + 1 < cols: 
-                candidates.append((row, col + 1)) 
-            candidates.append((row, col)) 
-            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: 
-                chosen = (row, col) 
-
-            r_c, c_c = chosen 
-            x_c = int((c_c + 1) * cell_w_px - cell_w_px * 0.1) 
-            y_c = int((r_c + 0.5) * cell_h_px) 
-
-            font_scale = max(0.6, min(1.6, cell_h_px / 60.0)) 
-            thickness = max(1, int(font_scale * 2)) 
+            # Position marks to the right of the answer, with fallback to left
+            x_c = int((col + 1) * cell_w_px - cell_w_px / 4) 
+            y_c = int((row + 0.5) * cell_h_px) 
+
+            # Use larger, more visible font
+            font_scale = max(1.0, min(2.0, cell_h_px / 40.0)) 
+            thickness = max(2, int(font_scale * 2)) 
             cv2.putText(img_cv, marks_text, (x_c, y_c), cv2.FONT_HERSHEY_SIMPLEX, 
                         font_scale, (0, 0, 255), thickness, cv2.LINE_AA) 
+            print(f"🖊 Marks annotated for page {page_num}, question {qid}: {marks_text}")
 
-        annotated_path = f"annotated_page_{p_index+1}.png" 
+        annotated_path = f"annotated_page_{page_num}.png" 
         cv2.imwrite(annotated_path, img_cv) 
         annotated_page_paths.append(annotated_path) 
         print("✅ Annotated page saved:", annotated_path) 
 
+    # Merge annotated pages into final PDF
+    print("📑 Merging annotated pages into final PDF...")
     with open(output_pdf, "wb") as f: 
         f.write(img2pdf.convert(annotated_page_paths)) 
 
@@ -649,8 +638,8 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
         return f"❌ Error: {e}", None, None, None, None
 
 # ---------------- GRADIO UI ---------------- 
-with gr.Blocks(title="LeadIB AI Grading (Final Flow — Verbose)") as demo: 
-    gr.Markdown("## 📘 LeadIB AI Grading — Final Flow\nUpload **Question Paper**, **Markscheme**, and **Student Answer Sheet**.\nFlow: merge QP+MS -> transcribe QP+MS (questions first, full markscheme) -> extract IDs -> transcribe AS with expected IDs -> grade -> (optional) imprint. Console prints show progress.") 
+with gr.Blocks(title=" AI Grading (Final Flow )") as demo: 
+    gr.Markdown("## 📘  AI Grading — Final Flow") 
 
     with gr.Row(): 
         qp_file = gr.File(label="📄 Upload Question Paper (PDF)")