Update main.py

main.py

@@ -1,10 +1,12 @@
 import os
+# Force stable engine to prevent PIR/DoubleAttribute crashes on CPU
 os.environ['FLAGS_enable_pir_api'] = '0'
 os.environ['FLAGS_use_mkldnn'] = '0'
 
 import re
 import cv2
 import json
+import uvicorn
 import numpy as np
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.responses import Response
@@ -21,28 +23,26 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ✅ LAZY LOAD OCR (IMPORTANT)
-ocr = None
+# Initialize PaddleOCR (More accurate for IDs than EasyOCR)
+ocr = PaddleOCR(use_angle_cls=True, lang='en', use_gpu=False, show_log=False)
 
-def get_ocr():
-    global ocr
-    if ocr is None:
-        ocr = PaddleOCR(use_angle_cls=True, lang='en', use_gpu=False)
-    return ocr
+def enhance(img):
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
+    return clahe.apply(gray)
 
 def clean_digits(text):
-    if not text:
-        return ""
-    s = str(text).upper()
-    s = s.replace('O','0').replace('D','0').replace('I','1').replace('L','1').replace('S','5').replace('B','8')
+    if not text: return ""
+    s = str(text).upper().replace('O','0').replace('D','0').replace('I','1').replace('L','1').replace('S','5').replace('B','8')
     return re.sub(r'[^0-9]', '', s)
 
-def mask_region(img, bbox):
+def mask_region(img, bbox, ratio=0.68):
     try:
+        # PaddleOCR bboxes: [[x1,y1], [x2,y1], [x2,y2], [x1,y2]]
         x1, y1 = int(bbox[0][0]), int(bbox[0][1])
         x2, y2 = int(bbox[2][0]), int(bbox[2][1])
         width = x2 - x1
-        mask_w = int(width * 0.68)
+        mask_w = int(width * ratio)
         cv2.rectangle(img, (x1, y1), (x1 + mask_w, y2), (0, 0, 0), -1)
     except:
         pass
@@ -54,59 +54,53 @@ async def process(file: UploadFile = File(...)):
         contents = await file.read()
         nparr = np.frombuffer(contents, np.uint8)
         img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        if img is None: return Response(content=json.dumps({"error": "Invalid image"}), status_code=400)
 
-        if img is None:
-            raise HTTPException(status_code=400, detail="Invalid image")
-
-        # Resize
+        # Standardize size
         h, w = img.shape[:2]
         img = cv2.resize(img, (1200, int(h * (1200 / w))))
-
-        # OCR
-        ocr_engine = get_ocr()
-        result = ocr_engine.ocr(img)
-
-        if not result:
-            raise HTTPException(status_code=400, detail="No text detected")
-
-        extracted_data = {"aadhaar": None, "vid": None}
-
-        for line in result:
-            if not isinstance(line, list):
-                continue
-
-            for word_info in line:
-                try:
-                    bbox = word_info[0]
-                    text = word_info[1][0]
-                    conf = word_info[1][1]
-
-                    if conf < 0.5:
-                        continue
-
-                    clean = clean_digits(text)
-
-                    if len(clean) == 12 and not clean.startswith(('0', '1')):
-                        extracted_data["aadhaar"] = clean
-                        img = mask_region(img, bbox)
-
-                    elif len(clean) == 16:
-                        extracted_data["vid"] = clean
-                        img = mask_region(img, bbox)
-
-                except:
-                    continue
-
-        _, buffer = cv2.imencode('.jpg', img)
-
+        
+        # OCR Inference
+        result = ocr.ocr(img)
+        extracted = {"aadhaar": None, "vid": None}
+
+        if result and isinstance(result, list):
+            flat_list = []
+            for line in result:
+                if not line: continue
+                for res in line:
+                    flat_list.append({"bbox": res[0], "text": res[1][0], "conf": res[1][1]})
+
+            for i, item in enumerate(flat_list):
+                if item["conf"] < 0.45: continue
+                clean = clean_digits(item["text"])
+                
+                # 12-digit Aadhaar
+                if len(clean) == 12 and not clean.startswith(('0', '1')):
+                    extracted["aadhaar"] = clean
+                    img = mask_region(img, item["bbox"], 0.68)
+                
+                # 16-digit VID
+                elif len(clean) == 16:
+                    extracted["vid"] = clean
+                    img = mask_region(img, item["bbox"], 0.75)
+                
+                # Split Aadhaar logic (e.g., 4444 5555 6666 across two boxes)
+                elif i + 1 < len(flat_list):
+                    combined = clean + clean_digits(flat_list[i+1]["text"])
+                    if len(combined) == 12:
+                        extracted["aadhaar"] = combined
+                        img = mask_region(img, item["bbox"], 1.0)
+                        img = mask_region(img, flat_list[i+1]["bbox"], 1.0)
+
+        _, 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_data),
-                "Access-Control-Expose-Headers": "x-data"
-            }
+            headers={"x-data": json.dumps(extracted), "Access-Control-Expose-Headers": "x-data"}
         )
-
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)