Update main.py

main.py

@@ -1,22 +1,22 @@
 import os
 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, File, HTTPException, Query
+from fastapi import FastAPI, UploadFile, Form, HTTPException
 from fastapi.responses import JSONResponse
-import uvicorn
 
 # -----------------------------
 # CONFIG
 # -----------------------------
 IMG_SIZE = (224, 224)
-MODEL_PATH = "final_model.keras"  # Updated to relative path for Docker
+MODEL_PATH = "./final_model.keras"  # Relative path for deployment
 class_names = ['Other', 'Aadhar Card', 'Pan Card', 'Voter Id']
 
-# Mapping for expected_type input (normalize to match class_names)
-EXPECTED_TYPE_MAPPING = {
+# Mapping for expected_type input to class names
+type_mapping = {
     "aadhar_card": "Aadhar Card",
     "pan_card": "Pan Card",
     "voter_id": "Voter Id"
@@ -111,42 +111,40 @@ def predict_file(path):
 # -----------------------------
 # FastAPI App
 # -----------------------------
-app = FastAPI(title="ID Document Validator", description="Predict document type and optionally validate against expected type.")
+app = FastAPI(title="ID Document Validator", description="Predict document type and authenticate based on 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.")
+    file: UploadFile,
+    expected_type: str = Form(None)  # Optional form field
 ):
+    # Validate file
+    if not file.filename:
+        raise HTTPException(status_code=400, detail="No file provided")
+    
     # Save uploaded file temporarily
-    temp_path = f"/tmp/{file.filename}"
-    with open(temp_path, "wb") as f:
-        f.write(await file.read())
+    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as temp_file:
+        temp_path = temp_file.name
+        temp_file.write(await file.read())
     
     try:
-        # Get prediction result
+        # Predict
         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
+        # Add authentication if expected_type is provided and no error
+        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"
         
-        return result
+        return JSONResponse(content=result)
     finally:
         # Clean up temp file
         if os.path.exists(temp_path):
-            os.remove(temp_path)
+            os.unlink(temp_path)
 
-# For local testing
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+# Health check endpoint
+@app.get("/")
+def read_root():
+    return {"message": "ID Validator API is running"}