Update OCR.py

OCR.py

@@ -5,45 +5,27 @@ from PIL import Image
 import io
 import re
 HF_TOKEN = os.environ.get("HF_TOKEN")
-# Set environment variable
 os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
 
-# Model and device setup
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 model_id = "google/paligemma-3b-mix-224"
-
-# Load model and processor
 model = PaliGemmaForConditionalGeneration.from_pretrained(model_id).to(device)
 processor = AutoProcessor.from_pretrained(model_id)
 
 def extract_text_from_image(image_content):
     image = Image.open(io.BytesIO(image_content))
-
-    # Prompt for detecting text
     prompt = "Extract all relevant details from this invoice."
-
-    # Prepare inputs for the model
     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
     input_len = inputs["input_ids"].shape[-1]
-
     with torch.inference_mode():
-        # Generate the output
         generation = model.generate(**inputs, max_new_tokens=100, do_sample=False)
         generation = generation[0][input_len:]
         decoded = processor.decode(generation, skip_special_tokens=True)
-
     return decoded
 
-def extract_text_from_pdf(pdf_content):
-    # For simplicity, let's assume you're converting the PDF to images first
-    # You may use libraries like pdf2image to convert PDF pages to images
-    # Then call extract_text_from_image for each image
-    pass
 
 def extract_invoice_details(text):
-    # Implement your logic to extract invoice details from the text
     details = {}
-    # Example extraction logic
     details['Invoice Number'] = re.search(r'Invoice Number: (\S+)', text).group(1) if re.search(r'Invoice Number: (\S+)', text) else 'N/A'
     details['Amount'] = re.search(r'Total Amount Due: (\S+)', text).group(1) if re.search(r'Total Amount Due: (\S+)', text) else 'N/A'
     details['Invoice Date'] = re.search(r'Invoice Date: (\S+)', text).group(1) if re.search(r'Invoice Date: (\S+)', text) else 'N/A'