Update main.py

main.py

@@ -183,101 +183,110 @@ def load_answer_key(pdf_bytes: bytes) -> dict:
 # FastAPI Endpoints
 ##############################################################
 
+from typing import Optional
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.responses import JSONResponse
+import numpy as np
+import cv2
+import json
+from PIL import Image
+
+app = FastAPI()
+
 @app.post("/process")
 async def process_pdfs(
-    student_pdf: UploadFile = File(..., description="PDF with all student answer sheets (one page per student)"),
-    paper_a_pdf: UploadFile = File(..., description="Answer key PDF for Paper A"),
-    paper_b_pdf: UploadFile = File(..., description="Answer key PDF for Paper B"),
-    paper_k_pdf: UploadFile = File(..., description="Answer key PDF for Paper K")
+    student_pdf: UploadFile = File(
+        ..., 
+        description="PDF with all student answer sheets (one page per student)"
+    ),
+    paper_k_pdf: UploadFile = File(
+        ...,
+        description="Answer key PDF for Paper K"
+    ),
+    paper_a_pdf: Optional[UploadFile] = File(
+        None,
+        description="(Optional) Answer key PDF for Paper A"
+    ),
+    paper_b_pdf: Optional[UploadFile] = File(
+        None,
+        description="(Optional) Answer key PDF for Paper B"
+    ),
 ):
     try:
-        # Read file bytes
+        # 1. Read the student and Paper K files (required)
         student_pdf_bytes = await student_pdf.read()
-        paper_a_bytes = await paper_a_pdf.read()
-        paper_b_bytes = await paper_b_pdf.read()
-        paper_k_bytes = await paper_k_pdf.read()
-        
-        # Preload answer keys from the three PDFs
+        paper_k_bytes     = await paper_k_pdf.read()
+
+        # 2. Read optional answer keys if provided
+        paper_a_bytes = await paper_a_pdf.read() if paper_a_pdf else None
+        paper_b_bytes = await paper_b_pdf.read() if paper_b_pdf else None
+
+        # 3. Build the answer_keys dict dynamically
         answer_keys = {
-            "A": load_answer_key(paper_a_bytes),
-            "B": load_answer_key(paper_b_bytes),
             "K": load_answer_key(paper_k_bytes)
         }
-        
-        # Convert the student answer PDF to images (each page = one student)
+        if paper_a_bytes is not None:
+            answer_keys["A"] = load_answer_key(paper_a_bytes)
+        if paper_b_bytes is not None:
+            answer_keys["B"] = load_answer_key(paper_b_bytes)
+
+        # 4. Convert the student PDF to images
         student_images = convert_from_bytes(student_pdf_bytes)
         all_results = []
-        
-        # Loop over all student pages
+
         for idx, page in enumerate(student_images):
             print(f"Processing student page {idx+1}...")
-            
-            # Convert the PIL image to OpenCV format for masking
-            page_cv = np.array(page)
-            page_cv = cv2.cvtColor(page_cv, cv2.COLOR_RGB2BGR)
-            height, width = page_cv.shape[:2]
-            
-            ###########################################################
-            # 1. Extract Candidate Information Region
-            ###########################################################
-            candidate_mask = np.zeros((height, width), dtype="uint8")
-            candidate_margin_top = int(height * 0.10)
-            candidate_margin_bottom = int(height * 0.75)
-            cv2.rectangle(candidate_mask, (0, candidate_margin_top), (width, height - candidate_margin_bottom), 255, -1)
-            masked_candidate = cv2.bitwise_and(page_cv, page_cv, mask=candidate_mask)
-            coords = cv2.findNonZero(candidate_mask)
+
+            # — Candidate Info Extraction (as before) —
+            page_cv = cv2.cvtColor(np.array(page), cv2.COLOR_RGB2BGR)
+            h, w = page_cv.shape[:2]
+            mask = np.zeros((h, w), dtype="uint8")
+            top = int(h * 0.10)
+            bottom = int(h * 0.75)
+            cv2.rectangle(mask, (0, top), (w, h - bottom), 255, -1)
+            masked = cv2.bitwise_and(page_cv, page_cv, mask=mask)
+            coords = cv2.findNonZero(mask)
             if coords is None:
-                continue  # Skip page if no candidate region is found.
-            x, y, w, h = cv2.boundingRect(coords)
-            cropped_candidate = masked_candidate[y:y+h, x:x+w]
-            candidate_pil = Image.fromarray(cv2.cvtColor(cropped_candidate, cv2.COLOR_BGR2RGB))
-            
-            # Extract candidate info using GenAI.
-            candidate_info_response = parse_info(candidate_pil)
-            candidate_info = extract_json_from_output(candidate_info_response)
-            
-            # Determine the candidate's paper.
-            paper = ""
-            if candidate_info and "Candidate Info" in candidate_info:
-                paper = candidate_info["Candidate Info"].get("Paper", "").strip()
+                continue
+            x, y, mw, mh = cv2.boundingRect(coords)
+            cand_pil = Image.fromarray(
+                cv2.cvtColor(masked[y:y+mh, x:x+mw], cv2.COLOR_BGR2RGB)
+            )
+            info_resp = parse_info(cand_pil)
+            cand_info = extract_json_from_output(info_resp)
+
+            # Determine which paper this student sat
+            paper = cand_info.get("Candidate Info", {}).get("Paper", "").strip().upper()
             if not paper:
-                paper = parse_paper(candidate_info_response)
-            paper = paper.upper()
+                paper = parse_paper(info_resp).upper()
             print(f"Student {idx+1} Paper: {paper}")
-            
-            # Retrieve the appropriate answer key.
+
+            # Skip if we don't have a key for that paper
             if paper not in answer_keys or answer_keys[paper] is None:
-                print(f"Error: Invalid or missing answer key for paper '{paper}' for student {idx+1}. Skipping.")
+                print(f"Skipping: no answer key for paper '{paper}'")
                 continue
-            correct_answer_key = answer_keys[paper]
-            
-            ###########################################################
-            # 2. Extract Student Answers from the Entire Page
-            ###########################################################
-            student_answers_response = parse_all_answers(page)
-            student_answers = extract_json_from_output(student_answers_response)
-            
-            ###########################################################
-            # 3. Calculate the Result for this Student
-            ###########################################################
-            result = calculate_result(student_answers, correct_answer_key)
-            
-            # Compile the result for this student.
-            result_card = {
-                "Student Index": idx + 1,
-                "Candidate Info": candidate_info.get("Candidate Info", {}) if candidate_info else {},
-                "Student Answers": student_answers,
-                "Correct Answer Key": correct_answer_key,
-                "Result": result
-            }
-            all_results.append(result_card)
-        
-        # Write the results to a file in the temporary folder.
+            correct_key = answer_keys[paper]
+
+            # — Student Answers Extraction —
+            ans_resp = parse_all_answers(page)
+            stud_answers = extract_json_from_output(ans_resp)
+
+            # — Scoring —
+            result = calculate_result(stud_answers, correct_key)
+            all_results.append({
+                "Student Index":     idx + 1,
+                "Candidate Info":    cand_info.get("Candidate Info", {}),
+                "Student Answers":   stud_answers,
+                "Correct Answer Key": correct_key,
+                "Result":            result
+            })
+
+        # 5. Save & return
         with open(RESULT_FILE, "w", encoding="utf-8") as f:
             json.dump({"results": all_results}, f, indent=2)
-        
+
         return JSONResponse(content={"results": all_results})
-    
+
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))