Update main.py

main.py

@@ -1,56 +1,308 @@
-import os
-import shutil
-from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.responses import JSONResponse
-from models import EvaluationRequest
-from pdf_processor import extract_answers_from_pdf, evaluate_student
+from fastapi import FastAPI, UploadFile, File, HTTPException, Query
+from fastapi.responses import JSONResponse, StreamingResponse
+import uvicorn
+import io
+import json
+import numpy as np
+import cv2
+from PIL import Image
+from pdf2image import convert_from_bytes
 
-app = FastAPI()
+# Import the Google GenAI client libraries.
+from google import genai
+from google.genai import types
 
-# Directory to temporarily store uploaded files.
-UPLOAD_DIR = "uploads"
-if not os.path.exists(UPLOAD_DIR):
-    os.makedirs(UPLOAD_DIR)
+# Initialize the GenAI client with your API key.
+client = genai.Client(api_key="AIzaSyDDDHg9GWl6-9aq9Wo43GHfk2wcakhgwBQ")
 
-@app.get("/")
-async def root():
+app = FastAPI(title="Student Result Card API")
+
+# -----------------------------
+# Preprocessing Methods
+# -----------------------------
+def preprocess_candidate_info(image_cv):
     """
-    Root endpoint that provides a welcome message.
+    Preprocess the image to extract the candidate information region.
+    Region is defined by a mask covering the top-left portion.
     """
-    return {
-        "message": "Welcome to the PDF Processing API. Use '/extract/' to extract answers from a PDF or '/evaluate/' to calculate marks using pre-extracted answers."
-    }
+    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))
 
-@app.post("/extract/")
-async def extract_pdf(file: UploadFile = File(...)):
+def preprocess_free_response(image_cv):
     """
-    Endpoint to extract answers from a PDF file.
+    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):
+    """
+    Extracts a JSON object from a string containing extra text.
+    """
+    start = output_str.find('{')
+    end = output_str.rfind('}')
+    if start == -1 or end == -1:
+        return None
+    json_str = output_str[start:end+1]
     try:
-        file_location = os.path.join(UPLOAD_DIR, file.filename)
-        with open(file_location, "wb") as f:
-            shutil.copyfileobj(file.file, f)
-        
-        result = extract_answers_from_pdf(file_location)
-        return JSONResponse(content=result.model_dump())
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-    finally:
-        if os.path.exists(file_location):
-            os.remove(file_location)
+        return json.loads(json_str)
+    except json.JSONDecodeError:
+        return None
+
+def get_student_info(image_input):
+    """
+    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)
 
-@app.post("/evaluate/")
-async def evaluate(evaluation_request: EvaluationRequest):
+def get_mcq_answers(image_input):
     """
-    Endpoint to evaluate student answers.
-    Expects a JSON payload with the pre-extracted answer key and student answers.
+    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.
+    """
+    output_format = """
+Answer in the following JSON format. Do not write anything else:
+{
+  "Answers": {
+    "1": "<option or text>",
+    "2": "<option or text>",
+    "3": "<option or text>",
+    "4": "<option or text>",
+    "5": "<option or text>",
+    "6": "<option or text>",
+    "7": "<option or text>",
+    "8": "<option or text>",
+    "9": "<option or text>",
+    "10": "<option or text>",
+    "11": "<free-text answer>",
+    "12": "<free-text answer>",
+    "13": "<free-text answer>",
+    "14": "<free-text answer>",
+    "15": "<free-text answer>"
+  }
+}
+"""
+    prompt = f"""
+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:
+{output_format}
+"""
+    response = client.models.generate_content(model="gemini-2.0-flash", contents=[prompt, image_input])
+    return extract_json_from_output(response.text)
+
+# -----------------------------
+# Method to calculate result card
+# -----------------------------
+def calculate_result(student_info, student_mcq, student_free, correct_answers):
+    """
+    Compares student's answers with the correct answers, calculates marks and percentage,
+    and returns a result card in JSON.
+    """
+    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"])
+    
+    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:
+            marks += 1
+            detailed[q] = {"Student": student_ans, "Correct": correct_ans, "Result": "Correct"}
+        else:
+            detailed[q] = {"Student": student_ans, "Correct": correct_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,
+        "Detailed Results": detailed
+    }
+    return result_card
+
+# -----------------------------
+# API Endpoint to process PDFs and return student result cards
+# -----------------------------
+@app.post("/process")
+async def process_pdfs(
+    student_pdf: UploadFile = File(...),
+    answer_key_pdf: UploadFile = File(...),
+    download: bool = Query(False, description="Set to true to download result card list as a JSON file")
+):
     try:
-        evaluation = evaluate_student(evaluation_request.answer_key, evaluation_request.student)
-        return JSONResponse(content=evaluation.model_dump())
+        # Read student PDF bytes and convert to images
+        student_bytes = await student_pdf.read()
+        student_images = convert_from_bytes(student_bytes)
+        
+        # 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)
+        
+        student_result_cards = []
+        
+        # Process each student page.
+        for idx, page in enumerate(student_images):
+            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)
+            
+            student_info = get_student_info(student_info_image)
+            student_mcq  = get_mcq_answers(mcq_image)
+            student_free = get_free_response_answers(free_image)
+            
+            result_card = calculate_result(student_info, student_mcq, student_free, correct_answers)
+            result_card["Student Index"] = idx + 1
+            student_result_cards.append(result_card)
+        
+        if download:
+            # Create downloadable JSON file
+            json_bytes = json.dumps({"result_cards": student_result_cards}, indent=2).encode("utf-8")
+            return StreamingResponse(
+                io.BytesIO(json_bytes),
+                media_type="application/json",
+                headers={"Content-Disposition": "attachment; filename=result_cards.json"}
+            )
+        else:
+            return JSONResponse(content={"result_cards": student_result_cards})
+    
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
 if __name__ == "__main__":
-    import uvicorn
     uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)