Create main.py

main.py

@@ -0,0 +1,337 @@
+import os
+import re
+import json
+import tempfile
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.preprocessing import image
+from pdf2image import convert_from_path
+from fastapi import FastAPI, UploadFile, Form, HTTPException, Header, Depends
+from fastapi.responses import JSONResponse
+import pytesseract
+from PIL import Image
+import requests
+
+# -----------------------------
+# CONFIG
+# -----------------------------
+IMG_SIZE = (224, 224)
+MODEL_PATH = "./final_model.keras"
+class_names = ['Other', 'Aadhaar Card', 'Pan Card', 'Voter Id']
+
+type_mapping = {
+    "aadhaar_card": "Aadhaar Card",
+    "pan_card": "Pan Card",
+    "voter_id": "Voter Id"
+}
+
+# Document types that support OCR extraction
+EXTRACTABLE_TYPES = {"Aadhaar Card", "Pan Card", "Voter Id"}
+
+# -----------------------------
+# AUTH CONFIG (HF SECRETS)
+# -----------------------------
+MASTER_SECRET_KEY = os.getenv("MASTER_SECRET_KEY")
+SARVAM_API_KEY = os.getenv("SARVAM_API_KEY")
+
+PROJECT_KEYS_JSON = os.getenv("PROJECT_KEYS_JSON", "{}")
+try:
+    PROJECT_KEYS = json.loads(PROJECT_KEYS_JSON)
+except Exception:
+    PROJECT_KEYS = {}
+
+# -----------------------------
+# AUTH VALIDATION
+# -----------------------------
+def verify_keys(
+    x_secret_key: str = Header(None),
+    x_project_id: str = Header(None),
+    x_project_key: str = Header(None),
+):
+    if not MASTER_SECRET_KEY:
+        raise HTTPException(status_code=500, detail="MASTER_SECRET_KEY not configured in Space secrets")
+
+    if not x_secret_key or x_secret_key != MASTER_SECRET_KEY:
+        raise HTTPException(status_code=401, detail="Invalid Secret Key")
+
+    if not x_project_id or not x_project_key:
+        raise HTTPException(status_code=401, detail="Project Id and Project Key are required")
+
+    if x_project_id not in PROJECT_KEYS:
+        raise HTTPException(status_code=401, detail=f"Unknown project: {x_project_id}")
+
+    if PROJECT_KEYS.get(x_project_id) != x_project_key:
+        raise HTTPException(status_code=401, detail="Invalid Project Key")
+
+    return True
+
+# -----------------------------
+# LOAD MODEL
+# -----------------------------
+model = tf.keras.models.load_model(MODEL_PATH)
+
+# -----------------------------
+# Predict a single image array (H,W,C)
+# -----------------------------
+def predict_array(img_array):
+    img_array = tf.cast(img_array, tf.float32)
+    img_array = tf.image.resize(img_array, IMG_SIZE)
+    img_array = tf.expand_dims(img_array, axis=0)
+
+    pred = model.predict(img_array, verbose=0)[0]
+    class_id = int(np.argmax(pred))
+    conf = float(np.max(pred))
+
+    return class_names[class_id], conf, pred
+
+# -----------------------------
+# Predict from image path
+# -----------------------------
+def predict_image(img_path):
+    img = image.load_img(img_path, target_size=IMG_SIZE, color_mode="rgb")
+    img_array = image.img_to_array(img)
+    label, conf, _ = predict_array(img_array)
+    return {
+        "type": label,
+        "confidence": round(conf, 6)
+    }
+
+# -----------------------------
+# Predict from PDF (max 2 pages only)
+# -----------------------------
+def predict_pdf(pdf_path, dpi=200, max_pages=2):
+    pages = convert_from_path(pdf_path, dpi=dpi)
+
+    if len(pages) > max_pages:
+        return {
+            "error": "Maximum page limit reached",
+            "max_pages": max_pages,
+            "found_pages": len(pages)
+        }
+
+    results = []
+    labels = []
+
+    for page in pages:
+        page_np = np.array(page)
+
+        if page_np.shape[-1] == 4:
+            page_np = page_np[..., :3]
+
+        label, conf, _ = predict_array(page_np)
+        results.append((label, conf))
+        labels.append(label)
+
+    if len(labels) == 2:
+        if labels[0] != labels[1]:
+            return {
+                "error": "Invalid input",
+                "reason": "Pages belong to different document types",
+                "page1_type": labels[0],
+                "page2_type": labels[1]
+            }
+
+        final_label = labels[0]
+        final_conf = float(max(results[0][1], results[1][1]))
+
+        return {
+            "type": final_label,
+            "confidence": round(final_conf, 6)
+        }
+
+    return {
+        "type": labels[0],
+        "confidence": round(float(results[0][1]), 6)
+    }
+
+# -----------------------------
+# Main predict dispatcher
+# -----------------------------
+def predict_file(path):
+    if not os.path.exists(path):
+        return {"error": "File not found", "path": path}
+
+    ext = os.path.splitext(path)[1].lower()
+
+    if ext in [".png", ".jpg", ".jpeg", ".bmp", ".webp"]:
+        return predict_image(path)
+
+    if ext == ".pdf":
+        return predict_pdf(path)
+
+    return {"error": "Unsupported file type", "ext": ext}
+
+# -----------------------------
+# OCR: Extract raw text from file
+# Uses eng+hin Tesseract langs to handle both English and Hindi text
+# -----------------------------
+def extract_text_from_file(path: str) -> str:
+    ext = os.path.splitext(path)[1].lower()
+    all_text = []
+
+    try:
+        if ext in [".png", ".jpg", ".jpeg", ".bmp", ".webp"]:
+            pil_img = Image.open(path).convert("RGB")
+            text = pytesseract.image_to_string(pil_img, lang="eng+hin")
+            all_text.append(text)
+
+        elif ext == ".pdf":
+            # Use higher DPI for better OCR accuracy on PDFs
+            pages = convert_from_path(path, dpi=250)
+            for page in pages:
+                text = pytesseract.image_to_string(page, lang="eng+hin")
+                all_text.append(text)
+
+    except Exception as e:
+        return f"OCR_ERROR: {str(e)}"
+
+    return "\n".join(all_text).strip()
+
+# -----------------------------
+# Sarvam AI: Extract structured fields from raw OCR text
+# -----------------------------
+def extract_fields_with_sarvam(raw_text: str, doc_type: str) -> dict:
+    empty_fields = {
+        "name": None,
+        "address": None,
+        "date_of_birth": None,
+        "gender": None,
+        "id_number": None,
+        "mobile_number": None,
+    }
+
+    if not SARVAM_API_KEY:
+        return {**empty_fields, "ocr_error": "SARVAM_API_KEY not configured in Space secrets"}
+
+    system_prompt = (
+        "You are a document field extractor for Indian identity documents. "
+        "Given raw OCR text, extract specific fields and return ONLY a valid JSON object. "
+        "No markdown, no explanation, no extra text — just the JSON.\n\n"
+        "Fields to extract:\n"
+        "  - name: Full name of the document holder\n"
+        "  - address: Full address (combine multi-line if needed)\n"
+        "  - date_of_birth: In DD/MM/YYYY or YYYY format; use year_of_birth key if only year is present\n"
+        "  - gender: Male / Female / Transgender\n"
+        "  - id_number: Aadhaar=12-digit number, PAN=10-char alphanumeric, Voter ID=alphanumeric card number\n"
+        "  - mobile_number: 10-digit mobile number if present\n\n"
+        "Use null for any field not found. Return a flat JSON object only."
+    )
+
+    user_prompt = (
+        f"Document Type: {doc_type}\n\n"
+        f"Raw OCR Text:\n{raw_text}\n\n"
+        "Extract the fields and return JSON."
+    )
+
+    try:
+        response = requests.post(
+            "https://api.sarvam.ai/v1/chat/completions",
+            headers={
+                "Authorization": f"Bearer {SARVAM_API_KEY}",
+                "Content-Type": "application/json"
+            },
+            json={
+                "model": "sarvam-m",
+                "messages": [
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt}
+                ],
+                "temperature": 0.0,
+                "max_tokens": 512
+            },
+            timeout=30
+        )
+
+        if response.status_code != 200:
+            return {
+                **empty_fields,
+                "sarvam_error": f"API returned {response.status_code}: {response.text[:300]}"
+            }
+
+        content = response.json()["choices"][0]["message"]["content"].strip()
+
+        # Strip markdown code fences if model wraps output
+        content = re.sub(r"^```(?:json)?\s*", "", content)
+        content = re.sub(r"\s*```$", "", content)
+
+        extracted = json.loads(content)
+
+        # Normalize keys — ensure all expected fields are present
+        fields = ["name", "address", "date_of_birth", "gender", "id_number", "mobile_number"]
+        normalized = {f: extracted.get(f) for f in fields}
+
+        # Handle year_of_birth fallback (if model returns it instead of date_of_birth)
+        if normalized["date_of_birth"] is None and "year_of_birth" in extracted:
+            normalized["date_of_birth"] = str(extracted["year_of_birth"])
+
+        return normalized
+
+    except json.JSONDecodeError as e:
+        return {**empty_fields, "sarvam_parse_error": f"JSON parse failed: {str(e)}"}
+    except requests.exceptions.Timeout:
+        return {**empty_fields, "sarvam_error": "Sarvam API request timed out"}
+    except Exception as e:
+        return {**empty_fields, "sarvam_error": str(e)}
+
+# -----------------------------
+# FastAPI App
+# -----------------------------
+app = FastAPI(
+    title="ID Document Validator",
+    description="Classify document type, authenticate, and extract structured fields via OCR + Sarvam AI."
+)
+
+@app.post("/predict")
+async def predict(
+    file: UploadFile,
+    expected_type: str = Form(None),
+    auth: bool = Depends(verify_keys)
+):
+    if not file.filename:
+        raise HTTPException(status_code=400, detail="No file provided")
+
+    suffix = os.path.splitext(file.filename)[1]
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
+        temp_path = temp_file.name
+        temp_file.write(await file.read())
+
+    try:
+        # Step 1: Classify document type
+        result = predict_file(temp_path)
+
+        # Step 2: Authenticate against expected_type
+        if expected_type and expected_type in type_mapping and "type" in result:
+            if result["type"] == type_mapping[expected_type]:
+                result["Authentication"] = "Valid"
+            else:
+                result["Authentication"] = "Not valid"
+
+        # Step 3: OCR + Sarvam field extraction (only for recognized ID types)
+        doc_type = result.get("type")
+        if doc_type in EXTRACTABLE_TYPES:
+            raw_text = extract_text_from_file(temp_path)
+
+            if raw_text.startswith("OCR_ERROR"):
+                result["extracted_fields"] = {
+                    "name": None,
+                    "address": None,
+                    "date_of_birth": None,
+                    "gender": None,
+                    "id_number": None,
+                    "mobile_number": None,
+                    "ocr_error": raw_text
+                }
+            else:
+                result["extracted_fields"] = extract_fields_with_sarvam(raw_text, doc_type)
+
+        return JSONResponse(content=result)
+
+    finally:
+        if os.path.exists(temp_path):
+            os.unlink(temp_path)
+
+
+@app.get("/")
+def read_root():
+    return {"message": "ID Validator API is running"}