Bhuvanesh2405/aadhaar_api

1from fastapi import FastAPI, UploadFile, File2from pydantic import BaseModel3from fastai.learner import load_learner4from fastai.vision.core import PILImage5import io6import base647 8app = FastAPI(title="Aadhaar Card Detection API")9learn = load_learner("aadhaar_classifier.pkl")10 11class ImageBytes(BaseModel):12    image_bytes: list13 14class ImageBase64(BaseModel):15    image_base64: str16 17def make_prediction(img):18    pred_class, pred_idx, probs = learn.predict(img)19    confidence = float(probs[pred_idx])20    if confidence < 0.80:21        pred_class = "notaadhaar"22    return str(pred_class), confidence23 24@app.get("/")25def home():26    return {"message": "Aadhaar Card Detection API is running"}27 28@app.get("/health")29def health():30    return {"status": "OK"}31 32@app.post("/predict/file")33async def predict_file(file: UploadFile = File(...)):34    try:35        image_bytes = await file.read()36        img = PILImage.create(io.BytesIO(image_bytes))37        pred_class, confidence = make_prediction(img)38        return {39            "status": "success",40            "prediction": pred_class,41            "is_aadhaar_card": pred_class.lower() == "aadhaar",42            "confidence": confidence43        }44    except Exception as e:45        return {"status": "error", "message": str(e)}46 47@app.post("/predict/bytes")48async def predict_bytes(data: ImageBytes):49    try:50        image_bytes = bytes(data.image_bytes)51        img = PILImage.create(io.BytesIO(image_bytes))52        pred_class, confidence = make_prediction(img)53        return {54            "status": "success",55            "prediction": pred_class,56            "is_aadhaar_card": pred_class.lower() == "aadhaar",57            "confidence": confidence58        }59    except Exception as e:60        return {"status": "error", "message": str(e)}61 62@app.post("/predict/base64")63async def predict_base64(data: ImageBase64):64    try:65        b64 = data.image_base6466        if "," in b64:67            b64 = b64.split(",")[1]68        image_bytes = base64.b64decode(b64)69        img = PILImage.create(io.BytesIO(image_bytes))70        pred_class, confidence = make_prediction(img)71        return {72            "status": "success",73            "prediction": pred_class,74            "is_aadhaar_card": pred_class.lower() == "aadhaar",75            "confidence": confidence76        }77    except Exception as e:78        return {"status": "error", "message": str(e)}