Spaces:
Sleeping
Sleeping
File size: 4,630 Bytes
a3de355 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
import pytesseract
from PIL import Image
import cv2
import numpy as np
import re
import os
import glob
import io
class OcrPriceExtractor:
def __init__(self, lang='eng'):
self.lang = lang
self.custom_config = r'--oem 3 --psm 6'
def extract_price(self, text):
pattern = r'Rp\s*\d+(?:[.,]\d+)*'
prices = re.findall(pattern, text)
return prices
def crop_image(self, image):
height = image.shape[0]
crop_top = int(height * 0.2) # 20% from top
crop_bottom = int(height * 0.2) # 20% from bottom
# Crop the image
cropped = image[crop_top:height-crop_bottom, :]
# cv2.imshow("cropped", cropped)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
return cropped
def extract_text_from_bytes(self, image_bytes):
try:
# Convert bytes to numpy array
nparr = np.frombuffer(image_bytes, np.uint8)
image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if image is None:
raise Exception("Could not decode image bytes")
cropped_image = self.crop_image(image)
pil_image = Image.fromarray(cropped_image)
# Extract text using pytesseract
text = pytesseract.image_to_string(pil_image, config=self.custom_config, lang=self.lang)
text = text.strip()
# Extract prices from the text
prices = self.extract_price(text)
return text, prices
except Exception as e:
print(f"Error during OCR: {str(e)}")
return None, None
def extract_text_from_bytesio(self, bytes_io):
try:
# Convert BytesIO to bytes
image_bytes = bytes_io.getvalue()
return self.extract_text_from_bytes(image_bytes)
except Exception as e:
print(f"Error during OCR from BytesIO: {str(e)}")
return None, None
def extract_text(self, image_input):
try:
if isinstance(image_input, io.BytesIO):
return self.extract_text_from_bytesio(image_input)
elif isinstance(image_input, bytes):
return self.extract_text_from_bytes(image_input)
image = cv2.imread(image_input)
if image is None:
raise Exception("Could not read image")
cropped_image = self.crop_image(image)
pil_image = Image.fromarray(cropped_image)
text = pytesseract.image_to_string(pil_image, config=self.custom_config, lang=self.lang)
text = text.strip()
prices = self.extract_price(text)
return text, prices
except Exception as e:
print(f"Error during OCR: {str(e)}")
return None, None
def test_all_images(self, directory="."):
image_extensions = ('*.jpg', '*.jpeg', '*.png', '*.bmp')
image_files = []
for ext in image_extensions:
image_files.extend(glob.glob(os.path.join(directory, ext)))
if not image_files:
print("No image files found in directory!\n")
return
for image_path in image_files:
print(f"\nProcessing: {image_path}")
print("-" * 50)
text, prices = self.extract_text(image_path)
self._print_results(text, prices)
with open(image_path, 'rb') as f:
bytes_io = io.BytesIO(f.read())
text, prices = self.extract_text(bytes_io)
self._print_results(text, prices)
print("-" * 50)
def _print_results(self, text, prices):
if text:
if prices:
print("\nTest Passed. Found Prices:")
for price in prices:
print(price)
else:
print("\nNo prices found in this image.")
else:
print("No text was extracted or an error occurred.")
if __name__ == "__main__":
import sys
ocr = OcrPriceExtractor()
if len(sys.argv) > 1 and sys.argv[1] == "--all":
ocr.test_all_images()
else:
image_path = "3.jpg"
with open(image_path, 'rb') as f:
bytes_io = io.BytesIO(f.read())
text, prices = ocr.extract_text(bytes_io)
if text:
print("\nFound Prices:")
for price in prices:
print(price)
else:
print("No text was extracted or an error occurred.")
|