Spaces:
Sleeping
Sleeping
| import json | |
| import os | |
| import re | |
| import secrets | |
| import cv2 | |
| import easyocr | |
| import numpy as np | |
| import uvicorn | |
| from fastapi import FastAPI, File, Header, HTTPException, UploadFile | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import Response | |
| APP_ENV = os.getenv("APP_ENV", "dev").lower() | |
| APP_PORT = int(os.getenv("PORT", "7860")) | |
| MAX_UPLOAD_MB = int(os.getenv("MAX_UPLOAD_MB", "10")) | |
| MAX_UPLOAD_BYTES = MAX_UPLOAD_MB * 1024 * 1024 | |
| API_KEY = os.getenv("MASKING_API_KEY", "") | |
| EXPOSE_DOCS = os.getenv("EXPOSE_DOCS", "true" if APP_ENV != "prod" else "false").lower() == "true" | |
| ALLOW_ORIGINS = [origin.strip() for origin in os.getenv("ALLOWED_ORIGINS", "").split(",") if origin.strip()] | |
| app = FastAPI( | |
| title="Aadhaar Strict Masking API", | |
| docs_url="/docs" if EXPOSE_DOCS else None, | |
| redoc_url="/redoc" if EXPOSE_DOCS else None, | |
| openapi_url="/openapi.json" if EXPOSE_DOCS else None, | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=ALLOW_ORIGINS or ["*"], | |
| allow_credentials=False, | |
| allow_methods=["POST", "GET"], | |
| allow_headers=["*"], | |
| ) | |
| reader = easyocr.Reader(["en"], gpu=False) | |
| def enhance_mobile_photo(img): | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| return cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(gray) | |
| def get_clean_digits(text): | |
| text = text.upper().replace("O", "0").replace("I", "1").replace("L", "1").replace("B", "8") | |
| return re.sub(r"[^0-9]", "", text) | |
| def mask_numeric_value(value, visible_last=4, mask_char="x"): | |
| if not value or len(value) <= visible_last: | |
| return value | |
| return f'{mask_char * (len(value) - visible_last)}{value[-visible_last:]}' | |
| def apply_mask(img, bbox, ratio): | |
| p1 = tuple(map(int, bbox[0])) | |
| p3 = tuple(map(int, bbox[2])) | |
| width = max(1, p3[0] - p1[0]) | |
| mask_width = int(width * ratio) | |
| cv2.rectangle(img, p1, (p1[0] + mask_width, p3[1]), (0, 0, 0), -1) | |
| return img | |
| def validate_api_key(x_api_key): | |
| if not API_KEY: | |
| return | |
| if not x_api_key or not secrets.compare_digest(x_api_key, API_KEY): | |
| raise HTTPException(status_code=401, detail="Unauthorized") | |
| async def health_check(): | |
| return {"status": "ok", "env": APP_ENV} | |
| async def process_document( | |
| file: UploadFile = File(...), | |
| x_api_key: str | None = Header(default=None, alias="x-api-key"), | |
| ): | |
| try: | |
| validate_api_key(x_api_key) | |
| contents = await file.read() | |
| if not contents: | |
| raise HTTPException(status_code=400, detail="Empty upload") | |
| if len(contents) > MAX_UPLOAD_BYTES: | |
| raise HTTPException(status_code=413, detail=f"File too large. Max allowed is {MAX_UPLOAD_MB} MB") | |
| img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR) | |
| if img is None: | |
| raise HTTPException(status_code=400, detail="Invalid image format") | |
| h, w = img.shape[:2] | |
| img = cv2.resize(img, (1200, int(h * (1200 / w)))) | |
| processed_img = enhance_mobile_photo(img) | |
| results = reader.readtext(processed_img, detail=1) | |
| aadhaar_values = [] | |
| vid_values = [] | |
| masked_aadhaar_keys = set() | |
| masked_vid_keys = set() | |
| for i, (bbox, text, conf) in enumerate(results): | |
| if conf < 0.35: | |
| continue | |
| clean = get_clean_digits(text) | |
| text_upper = text.upper() | |
| if len(clean) == 12 and not clean.startswith(("0", "1")): | |
| key = (clean, tuple(map(int, bbox[0])), tuple(map(int, bbox[2]))) | |
| if key not in masked_aadhaar_keys: | |
| aadhaar_values.append(clean) | |
| img = apply_mask(img, bbox, 0.68) | |
| masked_aadhaar_keys.add(key) | |
| continue | |
| if "VID" in text_upper and len(clean) >= 12: | |
| key = (clean[:16], tuple(map(int, bbox[0])), tuple(map(int, bbox[2]))) | |
| if key not in masked_vid_keys: | |
| vid_values.append(clean[:16]) | |
| img = apply_mask(img, bbox, 0.79) | |
| masked_vid_keys.add(key) | |
| elif "VID" in text_upper and i + 1 < len(results): | |
| next_bbox, next_text, next_conf = results[i + 1] | |
| if next_conf < 0.35: | |
| continue | |
| next_clean = get_clean_digits(next_text) | |
| if len(next_clean) >= 12: | |
| key = (next_clean[:16], tuple(map(int, next_bbox[0])), tuple(map(int, next_bbox[2]))) | |
| if key not in masked_vid_keys: | |
| vid_values.append(next_clean[:16]) | |
| img = apply_mask(img, next_bbox, 0.79) | |
| masked_vid_keys.add(key) | |
| elif len(clean) == 16: | |
| key = (clean, tuple(map(int, bbox[0])), tuple(map(int, bbox[2]))) | |
| if key not in masked_vid_keys: | |
| vid_values.append(clean) | |
| img = apply_mask(img, bbox, 0.79) | |
| masked_vid_keys.add(key) | |
| extracted = { | |
| "aadhaar": [mask_numeric_value(value, visible_last=4) for value in aadhaar_values] or None, | |
| "vid": [mask_numeric_value(value, visible_last=4) for value in vid_values] or None, | |
| } | |
| if isinstance(extracted["aadhaar"], list) and len(extracted["aadhaar"]) == 1: | |
| extracted["aadhaar"] = extracted["aadhaar"][0] | |
| if isinstance(extracted["vid"], list) and len(extracted["vid"]) == 1: | |
| extracted["vid"] = extracted["vid"][0] | |
| _, buffer = cv2.imencode(".jpg", img) | |
| return Response( | |
| content=buffer.tobytes(), | |
| media_type="image/jpeg", | |
| headers={ | |
| "x-data": json.dumps(extracted), | |
| "Access-Control-Expose-Headers": "x-data", | |
| "Cache-Control": "no-store", | |
| "Pragma": "no-cache", | |
| }, | |
| ) | |
| except Exception as e: | |
| if isinstance(e, HTTPException): | |
| raise e | |
| print("Error while processing document") | |
| raise HTTPException(status_code=500, detail="Internal server error") | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=APP_PORT) | |