Spaces:
Building
Building
| import os | |
| import re | |
| import cv2 | |
| import json | |
| import easyocr | |
| import uvicorn | |
| import numpy as np | |
| from fastapi import FastAPI, File, UploadFile, HTTPException | |
| from fastapi.responses import Response | |
| from fastapi.middleware.cors import CORSMiddleware | |
| app = FastAPI(title="Prohorizon Advanced Masking") | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # Initialize OCR once | |
| reader = easyocr.Reader(['en', 'hi'], gpu=False) | |
| def enhance_mobile_photo(img): | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)) | |
| enhanced = clahe.apply(gray) | |
| return enhanced | |
| def get_clean_digits(text): | |
| text = text.upper().replace('O', '0').replace('I', '1').replace('L', '1').replace('B', '8').replace('S', '5') | |
| return re.sub(r'[^0-9]', '', text) | |
| def apply_mask(img, bbox, mask_ratio): | |
| """Applies a black rectangle over a percentage of the detected text box.""" | |
| p1 = tuple(map(int, bbox[0])) # Top-left | |
| p3 = tuple(map(int, bbox[2])) # Bottom-right | |
| width = p3[0] - p1[0] | |
| mask_width = int(width * mask_ratio) | |
| # Draw the mask (0,0,0) is black | |
| cv2.rectangle(img, p1, (p1[0] + mask_width, p3[1]), (0, 0, 0), -1) | |
| return img | |
| async def process_document(file: UploadFile = File(...)): | |
| try: | |
| contents = await file.read() | |
| nparr = np.frombuffer(contents, np.uint8) | |
| img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) | |
| # 1. Resize for Speed & Accuracy | |
| h, w = img.shape[:2] | |
| img = cv2.resize(img, (1200, int(h * (1200/w)))) | |
| # 2. Enhance image for OCR pass | |
| processed_img = enhance_mobile_photo(img) | |
| # 3. OCR Pass | |
| results = reader.readtext(processed_img, detail=1) | |
| extracted = {"aadhaar": None, "vid": None, "name": None} | |
| # 4. Pattern Detection & Masking | |
| for i, (bbox, text, conf) in enumerate(results): | |
| clean = get_clean_digits(text) | |
| text_upper = text.upper() | |
| # --- 1. Aadhaar Masking (12 Digits) --- | |
| if 11 <= len(clean) <= 12: | |
| extracted["aadhaar"] = clean | |
| img = apply_mask(img, bbox, 0.68) # Masks 8 digits, leaves 4 | |
| continue # Move to next box once masked | |
| # --- 2. VID Masking (16 Digits) --- | |
| # Scenario A: VID label and numbers together or standalone 16 digits | |
| if len(clean) >= 15 or ( "VID" in text_upper and len(clean) >= 8): | |
| extracted["vid"] = clean | |
| img = apply_mask(img, bbox, 0.75) # Masks 12 digits, leaves 4 | |
| continue | |
| # Scenario B: "VID" label in current box, numbers in the NEXT box | |
| if "VID" in text_upper and i + 1 < len(results): | |
| next_bbox, next_text, _ = results[i+1] | |
| next_clean = get_clean_digits(next_text) | |
| if len(next_clean) >= 8: | |
| extracted["vid"] = next_clean | |
| img = apply_mask(img, next_bbox, 0.75) | |
| continue | |
| # --- 3. Name Extraction --- | |
| if not extracted["name"] and conf > 0.7: | |
| if not re.search(r'\d', text) and len(text.split()) >= 2: | |
| if not any(k in text_upper for k in ["GOVT", "INDIA", "MALE", "FEMALE", "DOB", "YEAR"]): | |
| extracted["name"] = text.strip() | |
| # 5. Final Image Response | |
| _, buffer = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), 95]) | |
| return Response( | |
| content=buffer.tobytes(), | |
| media_type="image/jpeg", | |
| headers={ | |
| "x-data": json.dumps(extracted), | |
| "Access-Control-Expose-Headers": "x-data" | |
| } | |
| ) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=8000) |