Update app.py

app.py

@@ -163,55 +163,7 @@ def create_model():
         return model 
     except Exception as e: 
         print("❌ Failed to create any Gemini model:", e) 
-        raise
-
-def wait_for_file_active(uploaded_file, timeout=300):
-    """
-    Wait for an uploaded file to become ACTIVE before proceeding.
-    Polls the file status every 2 seconds until ACTIVE or timeout.
-    """
-    import time
-    start_time = time.time()
-    
-    while time.time() - start_time < timeout:
-        try:
-            file_info = genai.get_file(uploaded_file.name)
-            state = file_info.state.name
-            print(f"📊 File {uploaded_file.name} status: {state}")
-            
-            if state == "ACTIVE":
-                print(f"✅ File {uploaded_file.name} is now ACTIVE and ready to use")
-                return True
-            elif state == "FAILED":
-                print(f"❌ File {uploaded_file.name} processing failed")
-                return False
-            
-            # Still processing, wait a bit
-            print(f"⏳ File still processing, waiting 2 seconds...")
-            time.sleep(2)
-            
-        except Exception as e:
-            print(f"❌ Error checking file status: {e}")
-            time.sleep(2)
-    
-    print(f"⏰ Timeout waiting for file {uploaded_file.name} to become active")
-    return False
-
-def validate_uploaded_file(uploaded_file):
-    """
-    Validate that an uploaded file is accessible and hasn't expired.
-    """
-    try:
-        # Try to access the file info to check if it's valid
-        file_info = genai.get_file(uploaded_file.name)
-        if file_info.state.name == "ACTIVE":
-            return True
-        else:
-            print(f"⚠️ File {uploaded_file.name} is not active (state: {file_info.state.name})")
-            return False
-    except Exception as e:
-        print(f"❌ File validation failed: {e}")
-        return False 
+        raise 
 
 def merge_pdfs(paths, output_path): 
     writer = PdfWriter() 
@@ -386,76 +338,71 @@ def extract_marks_from_grading(grading_text):
     return grading_json 
 
 # ---------------- MAPPING/IMPRINT HELPERS ---------------- 
-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.
+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}).
 For each question in the grading JSON, return the cell NUMBER where the FIRST STEP of that question begins.
-
-⚠ IMPORTANT RULES:
+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).
 - 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.
-- Each question should have unique cell number
+- 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}
-
+- 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 
 Return JSON only, like:
-[{{"page": 1, "question": "1(a)", "cell_number": 15}}, ...]
-
+[{{"question": "1.a", "cell_number": 15}}, ...]
 Grading JSON:
-{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}")
+{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) 
         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 batch sending to Gemini, 
-    then annotate and produce imprinted PDF using batch processing for better efficiency.
+    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. 
     """ 
     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 
@@ -463,7 +410,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", 20) 
+            num_font = ImageFont.truetype("arial.ttf", 16) 
         except Exception: 
             num_font = ImageFont.load_default() 
 
@@ -484,31 +431,34 @@ 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 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)}")
+    # 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 
 
     # 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 
 
-        # Filter mappings for this page
-        page_mappings = [m for m in all_mappings if m.get("page") == page_num]
-        
-        for item in page_mappings: 
+        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: 
@@ -524,24 +474,37 @@ def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, expec
             row = (cell_number - 1) // cols 
             col = (cell_number - 1) % cols 
 
-            # 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)) 
+            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)) 
             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_{page_num}.png" 
+        annotated_path = f"annotated_page_{p_index+1}.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)) 
 
@@ -590,36 +553,11 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
         merge_pdfs([qp_path, ms_path], merged_qpms_path)
         print("📎 Merged QP + MS ->", merged_qpms_path)
 
-        # Upload files to Gemini with error handling and validation
+        # Upload files to Gemini
         print("🔼 Uploading files to Gemini...")
-        try:
-            merged_uploaded = genai.upload_file(path=merged_qpms_path, display_name="QP+MS (merged)")
-            print("📤 QP+MS file uploaded, waiting for processing...")
-            if not wait_for_file_active(merged_uploaded):
-                raise Exception("QP+MS file failed to become active")
-            
-            ans_uploaded = genai.upload_file(path=ans_path, display_name="Answer Sheet")
-            print("📤 Answer Sheet uploaded, waiting for processing...")
-            if not wait_for_file_active(ans_uploaded):
-                raise Exception("Answer Sheet file failed to become active")
-                
-            print("✅ All files uploaded and active.")
-        except Exception as e:
-            print(f"❌ File upload/processing failed: {e}")
-            print("🔄 Retrying file upload...")
-            try:
-                merged_uploaded = genai.upload_file(path=merged_qpms_path, display_name="QP+MS (merged)")
-                if not wait_for_file_active(merged_uploaded):
-                    raise Exception("QP+MS retry failed to become active")
-                    
-                ans_uploaded = genai.upload_file(path=ans_path, display_name="Answer Sheet")
-                if not wait_for_file_active(ans_uploaded):
-                    raise Exception("Answer Sheet retry failed to become active")
-                    
-                print("✅ Retry upload successful.")
-            except Exception as retry_e:
-                print(f"❌ Retry failed: {retry_e}")
-                return f"❌ File upload error: {retry_e}", None, None, None, None
+        merged_uploaded = genai.upload_file(path=merged_qpms_path, display_name="QP+MS (merged)")
+        ans_uploaded = genai.upload_file(path=ans_path, display_name="Answer Sheet")
+        print("✅ Upload complete.")
 
         # Create model and print which selected
         model = create_model()