Update main.py

main.py

@@ -1,190 +1,58 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException, Query
-from fastapi.responses import JSONResponse, StreamingResponse
-import uvicorn
-import io
-import json
 import os
 import tempfile
+import io
+import json
 import numpy as np
 import cv2
 from PIL import Image
 from pdf2image import convert_from_bytes
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.responses import JSONResponse, StreamingResponse
+import uvicorn
 
+# Get API key from environment
 GENAI_API_KEY = os.getenv("GENAI_API_KEY")
 if not GENAI_API_KEY:
-    raise Exception("GENAI_API_KEY not set in .env file.")
+    raise Exception("GENAI_API_KEY not set in environment")
 
 # Import the Google GenAI client libraries.
 from google import genai
 from google.genai import types
 
-# Initialize the GenAI client with the API key from .env.
+# Initialize the GenAI client with the API key.
 client = genai.Client(api_key=GENAI_API_KEY)
 
 app = FastAPI(title="Student Result Card API")
 
-# Use the system temporary directory
+# Use system temporary directory to store the results file.
 TEMP_FOLDER = tempfile.gettempdir()
+RESULT_FILE = os.path.join(TEMP_FOLDER, "result_cards.json")
 
-# -----------------------------
-# Preprocessing Methods
-# -----------------------------
-def preprocess_candidate_info(image_cv):
-    """
-    Preprocess the image to extract the candidate information region.
-    Region is defined by a mask covering the top-left portion.
-    """
-    height, width = image_cv.shape[:2]
-    mask = np.zeros((height, width), dtype="uint8")
-    margin_top = int(height * 0.10)
-    margin_bottom = int(height * 0.25)
-    cv2.rectangle(mask, (0, margin_top), (width, height - margin_bottom), 255, -1)
-    masked = cv2.bitwise_and(image_cv, image_cv, mask=mask)
-    coords = cv2.findNonZero(mask)
-    x, y, w, h = cv2.boundingRect(coords)
-    cropped = masked[y:y+h, x:x+w]
-    return Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
-
-def preprocess_mcq(image_cv):
-    """
-    Preprocess the image to extract the MCQ answers region (questions 1 to 10).
-    Region is defined by a mask on the left side of the page.
-    """
-    height, width = image_cv.shape[:2]
-    mask = np.zeros((height, width), dtype="uint8")
-    margin_top = int(height * 0.27)
-    margin_bottom = int(height * 0.23)
-    right_boundary = int(width * 0.35)
-    cv2.rectangle(mask, (0, margin_top), (right_boundary, height - margin_bottom), 255, -1)
-    masked = cv2.bitwise_and(image_cv, image_cv, mask=mask)
-    coords = cv2.findNonZero(mask)
-    x, y, w, h = cv2.boundingRect(coords)
-    cropped = masked[y:y+h, x:x+w]
-    return Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
+##############################################################
+# Preprocessing & Extraction Functions
+##############################################################
 
-def preprocess_free_response(image_cv):
-    """
-    Preprocess the image to extract the free-response answers region (questions 11 to 15).
-    Region is defined by a mask on the middle-right part of the page.
-    """
-    height, width = image_cv.shape[:2]
-    mask = np.zeros((height, width), dtype="uint8")
-    margin_top = int(height * 0.27)
-    margin_bottom = int(height * 0.38)
-    left_boundary = int(width * 0.35)
-    right_boundary = int(width * 0.68)
-    cv2.rectangle(mask, (left_boundary, margin_top), (right_boundary, height - margin_bottom), 255, -1)
-    masked = cv2.bitwise_and(image_cv, image_cv, mask=mask)
-    coords = cv2.findNonZero(mask)
-    x, y, w, h = cv2.boundingRect(coords)
-    cropped = masked[y:y+h, x:x+w]
-    return Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
-
-def preprocess_full_answers(image_cv):
-    """
-    For extracting the correct answer key, we assume the entire page contains the answers.
-    """
-    return Image.fromarray(cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB))
-
-# -----------------------------
-# Extraction Methods using Gemini
-# -----------------------------
-def extract_json_from_output(output_str):
+def extract_json_from_output(output_str: str):
     """
     Extracts a JSON object from a string containing extra text.
     """
     start = output_str.find('{')
     end = output_str.rfind('}')
     if start == -1 or end == -1:
+        print("No JSON block found in the output.")
         return None
     json_str = output_str[start:end+1]
     try:
-        return json.loads(json_str)
-    except json.JSONDecodeError:
+        result = json.loads(json_str)
+        return result
+    except json.JSONDecodeError as e:
+        print("Error decoding JSON:", e)
         return None
 
-def get_student_info(image_input):
+def parse_all_answers(image_input: Image.Image) -> str:
     """
-    Extracts candidate information from an image.
-    """
-    output_format = """
-Answer in the following JSON format. Do not write anything else:
-{
-  "Candidate Info": {
-    "Name": "<name>",
-    "Number": "<number>",
-    "Country": "<country>",
-    "Level": "<level>"
-  }
-}
-"""
-    prompt = f"""
-You are an assistant that extracts candidate information from an image.
-The image contains details including name, candidate number, country, and level.
-Extract the information accurately and provide the result in JSON using the format below:
-{output_format}
-"""
-    response = client.models.generate_content(model="gemini-2.0-flash", contents=[prompt, image_input])
-    return extract_json_from_output(response.text)
-
-def get_mcq_answers(image_input):
-    """
-    Extracts multiple-choice answers (questions 1 to 10) from an image.
-    """
-    output_format = """
-Answer in the following JSON format do not write anything else:
-{
-  "Answers": {
-    "1": "<option>",
-    "2": "<option>",
-    "3": "<option>",
-    "4": "<option>",
-    "5": "<option>",
-    "6": "<option>",
-    "7": "<option>",
-    "8": "<option>",
-    "9": "<option>",
-    "10": "<option>"
-  }
-}
-"""
-    prompt = f"""
-You are an assistant that extracts MCQ answers from an image.
-The image is a screenshot of a 10-question multiple-choice answer sheet.
-Extract which option is marked for each question (1 to 10) and provide the answers in JSON using the format below:
-{output_format}
-"""
-    response = client.models.generate_content(model="gemini-2.0-flash", contents=[prompt, image_input])
-    return extract_json_from_output(response.text)
-
-def get_free_response_answers(image_input):
-    """
-    Extracts free-text answers (questions 11 to 15) from an image.
-    """
-    output_format = """
-Answer in the following JSON format. Do not write anything else:
-{
-  "Free Answers": {
-    "11": "<answer for question 11>",
-    "12": "<answer for question 12>",
-    "13": "<answer for question 13>",
-    "14": "<answer for question 14>",
-    "15": "<answer for question 15>"
-  }
-}
-"""
-    prompt = f"""
-You are an assistant that extracts free-text answers from an image.
-The image contains responses for questions 11 to 15.
-Extract the answers accurately and provide the result in JSON using the format below:
-{output_format}
-"""
-    response = client.models.generate_content(model="gemini-2.0-flash", contents=[prompt, image_input])
-    return extract_json_from_output(response.text)
-
-def get_all_answers(image_input):
-    """
-    Extracts all answers (questions 1 to 15) from an image of the correct answer key.
+    Extracts answers from an image of a 15-question answer sheet.
+    Returns the response text (JSON string).
     """
     output_format = """
 Answer in the following JSON format. Do not write anything else:
@@ -213,43 +81,83 @@ You are an assistant that extracts answers from an image.
 The image is a screenshot of an answer sheet containing 15 questions.
 For questions 1 to 10, the answers are multiple-choice selections.
 For questions 11 to 15, the answers are free-text responses.
-Extract the answer for each question and provide the result in JSON using the format below:
+Extract the answer for each question (1 to 15) and provide the result in JSON using the format below:
 {output_format}
 """
-    response = client.models.generate_content(model="gemini-2.0-flash", contents=[prompt, image_input])
-    return extract_json_from_output(response.text)
+    response = client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=[prompt, image_input]
+    )
+    return response.text
 
-# -----------------------------
-# Method to calculate result card
-# -----------------------------
-def calculate_result(student_info, student_mcq, student_free, correct_answers):
+def parse_info(image_input: Image.Image) -> str:
     """
-    Compares student's answers with the correct answers, calculates marks and percentage,
-    and returns a result card in JSON.
+    Extracts candidate information including name, number, country, level and paper from an image.
+    Returns the response text (JSON string).
     """
-    student_all = {}
-    if student_mcq and "Answers" in student_mcq:
-        student_all.update(student_mcq["Answers"])
-    if student_free and "Free Answers" in student_free:
-        student_all.update(student_free["Free Answers"])
-    
+    output_format = """
+Answer in the following JSON format. Do not write anything else:
+{
+  "Candidate Info": {
+    "Name": "<name>",
+    "Number": "<number>",
+    "Country": "<country>",
+    "Level": "<level>",
+    "Paper": "<paper>"
+  }
+}
+"""
+    prompt = f"""
+You are an assistant that extracts candidate information from an image.
+The image contains candidate details including name, candidate number, country, level and paper.
+Extract the information accurately and provide the result in JSON using the following format:
+{output_format}
+"""
+    response = client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=[prompt, image_input]
+    )
+    return response.text
+
+def parse_paper(student_info_text: str) -> str:
+    """
+    Extracts the Paper field from candidate information.
+    Returns the paper letter (e.g. "A", "B", or "K") as a string.
+    """
+    prompt = f"""
+You are an assistant that extracts the Paper from candidate information.
+The candidate information contains details including their paper designation.
+Extract the Paper value (one alphabet only) from the following:
+{student_info_text}
+"""
+    response = client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=[prompt, student_info_text]
+    )
+    return response.text.strip()
+
+def calculate_result(student_answers: dict, correct_answers: dict) -> dict:
+    """
+    Compares student's answers with the correct answers and calculates the score.
+    Assumes JSON structures with a top-level "Answers" key containing Q1 to Q15.
+    """
+    student_all = student_answers.get("Answers", {})
     correct_all = correct_answers.get("Answers", {})
     total_questions = 15
     marks = 0
     detailed = {}
     
     for q in map(str, range(1, total_questions + 1)):
-        student_ans = student_all.get(q, "").strip()
-        correct_ans = correct_all.get(q, "").strip()
-        if student_ans == correct_ans:
+        stud_ans = student_all.get(q, "").strip()
+        corr_ans = correct_all.get(q, "").strip()
+        if stud_ans == corr_ans:
             marks += 1
-            detailed[q] = {"Student": student_ans, "Correct": correct_ans, "Result": "Correct"}
+            detailed[q] = {"Student": stud_ans, "Correct": corr_ans, "Result": "Correct"}
         else:
-            detailed[q] = {"Student": student_ans, "Correct": correct_ans, "Result": "Incorrect"}
+            detailed[q] = {"Student": stud_ans, "Correct": corr_ans, "Result": "Incorrect"}
     
     percentage = (marks / total_questions) * 100
     result_card = {
-        "Candidate Info": student_info.get("Candidate Info", {}),
         "Total Marks": marks,
         "Total Questions": total_questions,
         "Percentage": percentage,
@@ -257,88 +165,140 @@ def calculate_result(student_info, student_mcq, student_free, correct_answers):
     }
     return result_card
 
-# -----------------------------
-# API Endpoint to process PDFs and return student result cards
-# -----------------------------
+##############################################################
+# Helper: Load and Process an Answer Key PDF (from bytes)
+##############################################################
+def load_answer_key(pdf_bytes: bytes) -> dict:
+    """
+    Converts a PDF (as bytes) to images, extracts the last page, and parses the answers.
+    Returns the parsed JSON answer key.
+    """
+    images = convert_from_bytes(pdf_bytes)
+    last_page_image = images[-1]
+    answer_key_response = parse_all_answers(last_page_image)
+    answer_key = extract_json_from_output(answer_key_response)
+    return answer_key
+
+##############################################################
+# FastAPI Endpoints
+##############################################################
+
 @app.post("/process")
 async def process_pdfs(
-    student_pdf: UploadFile = File(...),
-    answer_key_pdf: UploadFile = File(...),
-    download: bool = Query(True, description="Set to true to download result card list as a JSON file")
+    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")
 ):
     try:
-        # Read student PDF bytes and convert to images
-        student_bytes = await student_pdf.read()
-        student_images = convert_from_bytes(student_bytes)
+        # Read file bytes
+        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()
         
-        # Read answer key PDF bytes and convert to images; assume correct key is in the last page.
-        answer_key_bytes = await answer_key_pdf.read()
-        answer_key_images = convert_from_bytes(answer_key_bytes)
-        last_page = answer_key_images[-1]
-        last_page_cv = np.array(last_page)
-        last_page_cv = cv2.cvtColor(last_page_cv, cv2.COLOR_RGB2BGR)
-        correct_image = preprocess_full_answers(last_page_cv)
-        correct_answers = get_all_answers(correct_image)
+        # Preload answer keys from the three PDFs
+        answer_keys = {
+            "A": load_answer_key(paper_a_bytes),
+            "B": load_answer_key(paper_b_bytes),
+            "K": load_answer_key(paper_k_bytes)
+        }
         
-        student_result_cards = []
+        # Convert the student answer PDF to images (each page = one student)
+        student_images = convert_from_bytes(student_pdf_bytes)
+        all_results = []
         
-        # Process each student page.
+        # 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)
-            student_info_image = preprocess_candidate_info(page_cv)
-            mcq_image = preprocess_mcq(page_cv)
-            free_image = preprocess_free_response(page_cv)
+            height, width = page_cv.shape[:2]
             
-            student_info = get_student_info(student_info_image)
-            student_mcq  = get_mcq_answers(mcq_image)
-            student_free = get_free_response_answers(free_image)
+            ###########################################################
+            # 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)
+            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))
             
-            result_card = calculate_result(student_info, student_mcq, student_free, correct_answers)
-            result_card["Student Index"] = idx + 1
-            student_result_cards.append(result_card)
+            # 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()
+            if not paper:
+                paper = parse_paper(candidate_info_response)
+            paper = paper.upper()
+            print(f"Student {idx+1} Paper: {paper}")
+            
+            # Retrieve the appropriate answer key.
+            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.")
+                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)
         
-        response_data = {"result_cards": student_result_cards}
+        # Write the results to a file in the temporary folder.
+        with open(RESULT_FILE, "w", encoding="utf-8") as f:
+            json.dump({"results": all_results}, f, indent=2)
         
-        if download:
-            # Create downloadable JSON file and save to system temp folder
-            json_bytes = json.dumps(response_data, indent=2).encode("utf-8")
-            file_path = os.path.join(TEMP_FOLDER, "result_cards.json")
-            with open(file_path, "wb") as f:
-                f.write(json_bytes)
-            return StreamingResponse(
-                io.BytesIO(json_bytes),
-                media_type="application/json",
-                headers={"Content-Disposition": "attachment; filename=result_cards.json"}
-            )
-        else:
-            return JSONResponse(content=response_data)
+        return JSONResponse(content={"results": all_results})
     
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-# -----------------------------
-# New Download Endpoint
-# -----------------------------
 @app.get("/download")
-async def download_result_cards():
+async def download_results():
     """
-    Returns the previously generated result_cards.json file from the system temporary folder.
+    Returns the result JSON file stored in the temporary folder.
     """
-    file_path = os.path.join(TEMP_FOLDER, "result_cards.json")
-    if not os.path.exists(file_path):
-         raise HTTPException(status_code=404, detail="File not found")
+    if not os.path.exists(RESULT_FILE):
+        raise HTTPException(status_code=404, detail="Result file not found. Please run /process first.")
     return StreamingResponse(
-         open(file_path, "rb"),
-         media_type="application/json",
-         headers={"Content-Disposition": "attachment; filename=result_cards.json"}
+        open(RESULT_FILE, "rb"),
+        media_type="application/json",
+        headers={"Content-Disposition": f"attachment; filename=result_cards.json"}
     )
 
 @app.get("/")
 async def root():
     return {
         "message": "Welcome to the Student Result Card API.",
-        "usage": "POST PDFs to /process with 'student_pdf' and 'answer_key_pdf' fields. Use ?download=true for file download or GET /download to re-download the JSON file."
+        "usage": "POST PDFs to /process (student answer sheet, paper A, paper B, paper K). Then use /download to retrieve the results."
     }
 
 if __name__ == "__main__":