Update main.py

main.py

@@ -4,12 +4,12 @@ import json
 import easyocr
 import uvicorn
 import numpy as np
-from pyzbar import pyzbar # Dedicated QR Library
+from pyzbar import pyzbar
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.responses import Response
 from fastapi.middleware.cors import CORSMiddleware
 
-app = FastAPI(title="Aadhaar Strict Masking API")
+app = FastAPI(title="Aadhaar Masking HF Edition")
 
 app.add_middleware(
     CORSMiddleware,
@@ -19,8 +19,10 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Initialize EasyOCR (English only for stability)
+# Initialize once. Using English only to prevent memory crashes on HF Free Tier
+print("Loading OCR Engine...")
 reader = easyocr.Reader(['en'], gpu=False)
+print("Engine Loaded.")
 
 def get_clean_digits(text):
     text = text.upper().replace('O','0').replace('I','1').replace('L','1').replace('|','1')
@@ -31,35 +33,40 @@ def apply_mask(img, bbox, ratio=0.7):
     p3 = tuple(map(int, bbox[2]))
     width = p3[0] - p1[0]
     mask_width = int(width * ratio)
+    # Applying a black mask
     cv2.rectangle(img, p1, (p1[0] + mask_width, p3[1]), (0, 0, 0), -1)
     return img
 
 def mask_qr_with_zbar(img):
     """
-    Uses pyzbar to find the ACTUAL QR code pattern. 
-    This prevents the 'Full Black' image error.
+    Detects actual QR patterns. If pyzbar is missing system libs, 
+    this fails gracefully instead of crashing.
     """
-    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    # Detect QR codes
-    qr_codes = pyzbar.decode(gray)
-    
-    for qr in qr_codes:
-        (x, y, w, h) = qr.rect
-        # Mask the detected QR area with a small padding
-        cv2.rectangle(img, (x-5, y-5), (x + w + 5, y + h + 5), (0, 0, 0), -1)
-    
+    try:
+        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        qr_codes = pyzbar.decode(gray)
+        for qr in qr_codes:
+            (x, y, w, h) = qr.rect
+            # Mask with slight padding
+            cv2.rectangle(img, (x-10, y-10), (x+w+10, y+h+10), (0, 0, 0), -1)
+    except Exception as e:
+        print(f"QR detection skipped: {e}")
     return img
 
 @app.post("/v1/aadhaar/process")
 async def process_document(file: UploadFile = File(...)):
     try:
         contents = await file.read()
-        img = cv2.imdecode(np.frombuffer(contents, np.uint8), cv2.IMREAD_COLOR)
+        nparr = np.frombuffer(contents, np.uint8)
+        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+        if img is None:
+            raise HTTPException(status_code=400, detail="Invalid image file")
 
-        # 1. Mask QR Code using pyzbar (Safe and Precise)
+        # 1. Mask QR Code specifically (No more full-black image)
         img = mask_qr_with_zbar(img)
 
-        # 2. Resize
+        # 2. Resize for OCR (Standardize 1200px width)
         h, w = img.shape[:2]
         img = cv2.resize(img, (1200, int(h * (1200/w))))
 
@@ -71,29 +78,26 @@ async def process_document(file: UploadFile = File(...)):
             clean = get_clean_digits(text)
             text_upper = text.upper()
 
-            # --- AADHAAR (12 Digits) ---
+            # AADHAAR: 12 digits, must not start with 0 or 1
             if len(clean) == 12 and not clean.startswith(('0','1')):
                 extracted["aadhaar"] = clean
-                img = apply_mask(img, bbox, 0.7) 
+                img = apply_mask(img, bbox, 0.72)
                 continue
 
-            # --- VID (16 Digits) ---
-            # We only mask if it is 16 digits OR the word VID is right there.
-            if len(clean) == 16 or "VID" in text_upper:
-                target_bbox = bbox
-                target_clean = clean
+            # VID: 16 digits or text context
+            if "VID" in text_upper or len(clean) == 16:
+                target_bbox, target_clean = bbox, clean
                 
-                # Handle split VID label and number
+                # Check next block if VID label is separate from numbers
                 if len(clean) < 12 and i + 1 < len(results):
                     next_bbox, next_text, _ = results[i+1]
                     next_clean = get_clean_digits(next_text)
                     if len(next_clean) >= 12:
-                        target_bbox = next_bbox
-                        target_clean = next_clean
+                        target_bbox, target_clean = next_bbox, next_clean
                 
                 if len(target_clean) >= 12:
                     extracted["vid"] = target_clean
-                    img = apply_mask(img, target_bbox, 0.75)
+                    img = apply_mask(img, target_bbox, 0.78)
 
         _, buffer = cv2.imencode('.jpg', img)
         return Response(
@@ -106,7 +110,8 @@ async def process_document(file: UploadFile = File(...)):
         )
 
     except Exception as e:
+        print(f"Error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=7860) # Port 7860 is default for HF Spaces