Update main.py

main.py

@@ -1,48 +1,152 @@
-# app.py
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
-from fastapi.responses import JSONResponse
-from predictor import predict_file
-import shutil
 import os
-import uuid
-import tempfile
+import json
+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, File, HTTPException, Query
+from fastapi.responses import JSONResponse
+import uvicorn
 
-app = FastAPI(title="Indian ID Document Classifier API")
+# -----------------------------
+# CONFIG
+# -----------------------------
+IMG_SIZE = (224, 224)
+MODEL_PATH = "final_model (2).keras"  # Updated to relative path for Docker
+class_names = ['Other', 'Aadhar Card', 'Pan Card', 'Voter Id']
 
-ALLOWED_TYPES = {"Aadhar Card", "Pan Card", "Voter Id"}
+# Mapping for expected_type input (normalize to match class_names)
+EXPECTED_TYPE_MAPPING = {
+    "aadhar_card": "Aadhar Card",
+    "pan_card": "Pan Card",
+    "voter_id": "Voter Id"
+}
 
+# -----------------------------
+# LOAD MODEL
+# -----------------------------
+model = tf.keras.models.load_model(MODEL_PATH)
 
-@app.post("/predict")
-async def predict_document(
-    file: UploadFile = File(...),
-    expected_type: str | None = Form(None),
-):
-    with tempfile.TemporaryDirectory() as temp_dir:
-        file_ext = os.path.splitext(file.filename)[1].lower()
-        temp_path = os.path.join(temp_dir, f"upload_{uuid.uuid4()}{file_ext}")
+# -----------------------------
+# 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)
+    # DO NOT divide by 255.0 (training did not use Rescaling)
+    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
 
-        try:
-            with open(temp_path, "wb") as buffer:
-                shutil.copyfileobj(file.file, buffer)
+# -----------------------------
+# 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)
+    }
 
-            result = predict_file(temp_path)
+# -----------------------------
+# 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)
+        # RGBA -> RGB
+        if page_np.shape[-1] == 4:
+            page_np = page_np[..., :3]
+        label, conf, _ = predict_array(page_np)
+        results.append((label, conf))
+        labels.append(label)
+    # 2 pages validation
+    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]
+            }
+        # both same -> return single output
+        final_label = labels[0]
+        final_conf = float(max(results[0][1], results[1][1]))
+        return {
+            "type": final_label,
+            "confidence": round(final_conf, 6)
+        }
+    # 1-page pdf normal output
+    return {
+        "type": labels[0],
+        "confidence": round(float(results[0][1]), 6)
+    }
 
-            # Authentication field logic
-            if expected_type is not None:
-                expected = expected_type.strip()
-                if expected in ALLOWED_TYPES:
-                    predicted = result.get("type")
-                    if isinstance(predicted, str) and predicted == expected:
-                        result["Authentication"] = "Valid"
-                    else:
-                        result["Authentication"] = "Not valid"
+# -----------------------------
+# Main function
+# -----------------------------
+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}
 
-            return JSONResponse(content=result)
-
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Server error: {str(e)}")
+# -----------------------------
+# FastAPI App
+# -----------------------------
+app = FastAPI(title="ID Document Validator", description="Predict document type and optionally validate against expected type.")
 
+@app.post("/predict")
+async def predict(
+    file: UploadFile = File(...),
+    expected_type: str = Query(None, description="Optional: Expected document type ('aadhar_card', 'pan_card', 'voter_id'). If provided, adds 'Authentication' to response.")
+):
+    # Save uploaded file temporarily
+    temp_path = f"/tmp/{file.filename}"
+    with open(temp_path, "wb") as f:
+        f.write(await file.read())
+    
+    try:
+        # Get prediction result
+        result = predict_file(temp_path)
+        
+        # If there's an error in result, return it as-is
+        if "error" in result:
+            return JSONResponse(status_code=400, content=result)
+        
+        # If expected_type is provided and valid, add Authentication
+        if expected_type and expected_type in EXPECTED_TYPE_MAPPING:
+            expected_label = EXPECTED_TYPE_MAPPING[expected_type]
+            predicted_label = result["type"]
+            authentication = "Valid" if predicted_label == expected_label else "Not Valid"
+            result["Authentication"] = authentication
+        elif expected_type and expected_type not in EXPECTED_TYPE_MAPPING:
+            # Invalid expected_type, but still return prediction without Authentication
+            pass  # Just proceed without adding Authentication
+        
+        return result
+    finally:
+        # Clean up temp file
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
 
-@app.get("/health")
-async def health_check():
-    return {"status": "ok"}
+# For local testing
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)