Spaces:
Running
Running
| import easyocr | |
| import json | |
| import re | |
| import cv2 | |
| import easyocr | |
| import google.generativeai as genai | |
| langs = {'french': 'fr', 'english': 'en', 'arabic': 'ar'} | |
| GOOGLE_API_KEY = "AIzaSyC6NXwTrucSl2JkY23YWsucFZPMBDoaqJw" | |
| genai.configure(api_key=GOOGLE_API_KEY) | |
| model = genai.GenerativeModel('gemini-pro') | |
| def get_ocr(image,reader): | |
| result = reader.readtext(image) | |
| return result | |
| def get_input(result, sep = " "): | |
| a = {} | |
| for (bbox, text, prob) in result: | |
| if prob<0.3: | |
| continue | |
| k = True | |
| for row in a: | |
| if abs(bbox[0][1] - row)<=5: | |
| k = False | |
| a[row]+= sep | |
| a[row]+= text | |
| if k: | |
| a[bbox[0][1]] = text | |
| inputt = "" | |
| for row in a: | |
| inputt+= a[row] | |
| inputt+= "\n" | |
| return inputt | |
| def imp_ocr(result, image, reader): | |
| v_exp_rate = 0.2 | |
| h_exp_rate = 0.3 | |
| imp_result = [] | |
| for i in range(len(result)): | |
| prob = result[i][2] | |
| if prob < 0.1: | |
| continue | |
| if prob < 0.9: | |
| bbox = result[i][0] | |
| x = int(bbox[0][0] - (bbox[2][0] - bbox[0][0])*h_exp_rate/2) | |
| x_h = int(bbox[2][0] + (bbox[2][0] - bbox[0][0])*h_exp_rate/2) | |
| y = int(bbox[0][1] - (bbox[2][1] - bbox[0][1])*v_exp_rate/2) | |
| y_h = int(bbox[2][1] + (bbox[2][1] - bbox[0][1])*v_exp_rate/2) | |
| x = max(x, 0) | |
| x_h = min(x_h, image.shape[1]) | |
| y = max(y, 0) | |
| y_h = min(y_h, image.shape[0]) | |
| sub_img = image[y:y_h, x:x_h] | |
| res = get_ocr(sub_img,reader) | |
| if not res: | |
| imp_result.append(result[i]) | |
| continue | |
| if len(res)>1: | |
| res = sorted(res, key=lambda x: x[2], reverse=True) | |
| if res[0][2] >= prob: | |
| imp_result.append((result[i][0], res[0][1], res[0][2])) | |
| else: | |
| imp_result.append(result[i]) | |
| else: | |
| imp_result.append(result[i]) | |
| return imp_result | |
| def extract_data(lang, to_be_extracted, image): | |
| reader = easyocr.Reader([langs[lang]]) | |
| ocr_result = get_ocr(image,reader) | |
| imp_result = imp_ocr(ocr_result, image, reader) | |
| inputt = get_input(imp_result, sep = " ") | |
| return get_output(inputt, to_be_extracted, lang) | |
| def get_output(inputt, to_be_extracted, lang): | |
| prompt = f""" | |
| Bellow is the ouptut text of an OCR system. The text is in {lang}. | |
| Your job is: | |
| 1. Format the output in a python dictionary format with the following keys : | |
| {to_be_extracted}. | |
| 2. correct any mistakes such as date formating (should be dd/mm/yyyy) or spelling mistakes (this is important). | |
| here is your input: | |
| {inputt} | |
| """ | |
| response = model.generate_content(prompt) | |
| data = extract_json(response.text) | |
| return data | |
| def extract_json(text_response): | |
| # This pattern matches a string that starts with '{' and ends with '}' | |
| pattern = r'\{[^{}]*\}' | |
| matches = re.finditer(pattern, text_response) | |
| json_objects = [] | |
| for match in matches: | |
| json_str = match.group(0) | |
| try: | |
| # Validate if the extracted string is valid JSON | |
| json_obj = eval(json_str) | |
| json_objects.append(json_obj) | |
| except json.JSONDecodeError: | |
| # Extend the search for nested structures | |
| extended_json_str = extend_search(text_response, match.span()) | |
| try: | |
| json_obj = eval(extended_json_str) | |
| json_objects.append(json_obj) | |
| except json.JSONDecodeError: | |
| # Handle cases where the extraction is not valid JSON | |
| continue | |
| if json_objects: | |
| return json_objects[0] | |
| else: | |
| return {} | |
| def extend_search(text, span): | |
| # Extend the search to try to capture nested structures | |
| start, end = span | |
| nest_count = 0 | |
| for i in range(start, len(text)): | |
| if text[i] == '{': | |
| nest_count += 1 | |
| elif text[i] == '}': | |
| nest_count -= 1 | |
| if nest_count == 0: | |
| return text[start:i+1] | |
| return text[start:end] |