Update modules/kyc_processor.py

modules/kyc_processor.py

@@ -1,8 +1,16 @@
-import pytesseract
-import cv2
+from PIL import Image
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import torch
+
+# Load once at startup
+processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
+model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-stage1")
 
 def extract_text_from_id(image_path):
-    image = cv2.imread(image_path)
-    if image is None:
-        raise ValueError(f"Invalid image file: {image_path}")
-    return pytesseract.image_to_string(image)
+    image = Image.open(image_path).convert("RGB")
+
+    pixel_values = processor(images=image, return_tensors="pt").pixel_values
+    generated_ids = model.generate(pixel_values)
+
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return generated_text