Update OCR.py

OCR.py

@@ -14,7 +14,7 @@ processor = AutoProcessor.from_pretrained(model_id)
 
 def extract_text_from_image(image_content):
     image = Image.open(io.BytesIO(image_content))
-    prompt =  "Extract the following details from this invoice: Invoice Number, Amount, Invoice Date, Due Date."
+    prompt =  "Extract the following details from this invoice: Invoice Number, Total Amount, Invoice Date."
     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
     input_len = inputs["input_ids"].shape[-1]
     with torch.inference_mode():
@@ -29,5 +29,4 @@ def extract_invoice_details(text):
     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'
-    details['Due Date'] = re.search(r'Due Date: (\S+)', text).group(1) if re.search(r'Due Date: (\S+)', text) else 'N/A'
     return details