import re import time import base64 import json import cv2 import numpy as np from typing import List, Optional from openai import OpenAI import numpy as np from .base import register_OCR, OCRBase, TextBlock @register_OCR("llm_ocr") class LLM_OCR(OCRBase): lang_map = { "Auto Detect": None, "Afrikaans": "af", "Albanian": "sq", "Amharic": "am", "Arabic": "ar", "Armenian": "hy", "Assamese": "as", "Azerbaijani": "az", "Bangla": "bn", "Basque": "eu", "Belarusian": "be", "Bengali": "bn", "Bosnian": "bs", "Breton": "br", "Bulgarian": "bg", "Burmese": "my", "Catalan": "ca", "Cebuano": "ceb", "Cherokee": "chr", "Chinese (Simplified)": "zh-CN", "Chinese (Traditional)": "zh-TW", "Corsican": "co", "Croatian": "hr", "Czech": "cs", "Danish": "da", "Dutch": "nl", "English": "en", "Esperanto": "eo", "Estonian": "et", "Faroese": "fo", "Filipino": "fil", "Finnish": "fi", "French": "fr", "Frisian": "fy", "Galician": "gl", "Georgian": "ka", "German": "de", "Greek": "el", "Gujarati": "gu", "Haitian Creole": "ht", "Hausa": "ha", "Hawaiian": "haw", "Hebrew": "he", "Hindi": "hi", "Hmong": "hmn", "Hungarian": "hu", "Icelandic": "is", "Igbo": "ig", "Indonesian": "id", "Interlingua": "ia", "Irish": "ga", "Italian": "it", "Japanese": "ja", "Javanese": "jv", "Kannada": "kn", "Kazakh": "kk", "Khmer": "km", "Korean": "ko", "Kurdish": "ku", "Kyrgyz": "ky", "Lao": "lo", "Latin": "la", "Latvian": "lv", "Lithuanian": "lt", "Luxembourgish": "lb", "Macedonian": "mk", "Malagasy": "mg", "Malay": "ms", "Malayalam": "ml", "Maltese": "mt", "Maori": "mi", "Marathi": "mr", "Mongolian": "mn", "Nepali": "ne", "Norwegian": "no", "Occitan": "oc", "Oriya": "or", "Pashto": "ps", "Persian": "fa", "Polish": "pl", "Portuguese": "pt", "Punjabi": "pa", "Quechua": "qu", "Romanian": "ro", "Russian": "ru", "Samoan": "sm", "Scots Gaelic": "gd", "Serbian (Cyrillic)": "sr-Cyrl", "Serbian (Latin)": "sr-Latn", "Shona": "sn", "Sindhi": "sd", "Sinhala": "si", "Slovak": "sk", "Slovenian": "sl", "Somali": "so", "Spanish": "es", "Sundanese": "su", "Swahili": "sw", "Swedish": "sv", "Tagalog": "tl", "Tajik": "tg", "Tamil": "ta", "Tatar": "tt", "Telugu": "te", "Thai": "th", "Tibetan": "bo", "Tigrinya": "ti", "Tongan": "to", "Turkish": "tr", "Ukrainian": "uk", "Urdu": "ur", "Uyghur": "ug", "Uzbek": "uz", "Vietnamese": "vi", "Welsh": "cy", "Xhosa": "xh", "Yiddish": "yi", "Yoruba": "yo", "Zulu": "zu", } popular_models = [ "OAI: gpt-4-vision-preview", "OAI: gpt-4", "OAI: gpt-3.5-turbo", "GGL: gemini-1.5-pro-latest", "GGL: gemini-2.0-flash-exp", "GGL: gemini-2.0-flash", ] params = { "provider": { "type": "selector", "options": ["OpenAI", "Google"], "value": "OpenAI", "description": "Select the LLM provider.", }, "api_key": {"value": "", "description": "Your API key."}, "endpoint": { "value": "", # Default to empty, allowing provider to dictate "description": "Base URL for the API. Leave empty to use provider default.", }, "model": { "type": "selector", "options": popular_models, "value": "", # Default to empty, allowing provider to dictate "description": "Select the model to use. Leave empty to use provider default. (Provider prefix indicates the provider).", }, "override_model": { "value": "", "description": "Specify a custom model name to override the selected model.", }, "language": { "type": "selector", "options": list(lang_map.keys()), "value": "Auto Detect", "description": "Language for OCR.", }, "prompt": { "value": "Recognize the text in this image.", "description": "Default prompt for OCR.", }, "system_prompt": { "type": "editor", "value": "", "description": "Optional system prompt to guide the model's behavior.", }, "proxy": { "value": "", "description": "Proxy address (e.g., http(s)://user:password@host:port or socks4/5://user:password@host:port)", }, "delay": {"value": 1.0, "description": "Delay in seconds between requests."}, "requests_per_minute": { "value": 15, "description": "Maximum number of requests per minute (0 for no limit).", }, "description": "OCR using various LLMs compatible with the OpenAI API.", } def __init__(self, **params) -> None: super().__init__(**params) self.last_request_time = 0 self.client = None self._initialize_client() self.request_count_minute = 0 self.minute_start_time = time.time() def _initialize_client(self): import httpx # Configure proxies using mounts if self.proxy: proxy_mounts = { "http://": httpx.HTTPTransport(proxy=self.proxy), "https://": httpx.HTTPTransport(proxy=self.proxy), } transport = httpx.Client(mounts=proxy_mounts) else: transport = httpx.Client() # No proxy # Determine the endpoint endpoint = self.endpoint if not endpoint: # If endpoint is empty, use provider default provider = self.provider if provider == "OpenAI": endpoint = "https://api.openai.com/v1" elif provider == "Google": endpoint = "https://generativelanguage.googleapis.com/v1beta/openai" else: endpoint = "https://api.openai.com/v1" # Default self.client = OpenAI( api_key=self.api_key, base_url=endpoint, http_client=transport ) @property def provider(self): return self.get_param_value("provider") @property def request_delay(self): try: return float(self.get_param_value("delay")) except (ValueError, TypeError): return 1.0 @property def api_key(self): return self.get_param_value("api_key") @property def endpoint(self): return self.get_param_value("endpoint") @property def model(self): return self.get_param_value("model") @property def override_model(self): return self.get_param_value("override_model") @property def language(self): lang_name = self.get_param_value("language") return self.lang_map.get(lang_name) @property def prompt(self): return self.get_param_value("prompt") @property def system_prompt(self): return self.get_param_value("system_prompt") @property def proxy(self): return self.get_param_value("proxy") @property def requests_per_minute(self): return int(self.get_param_value("requests_per_minute")) def _respect_delay(self): current_time = time.time() # Handle RPM limit if self.requests_per_minute > 0: if current_time - self.minute_start_time >= 60: self.request_count_minute = 0 self.minute_start_time = current_time if self.request_count_minute >= self.requests_per_minute: wait_time = 62 - (current_time - self.minute_start_time) if wait_time > 0: if self.debug_mode: self.logger.info( f"Reached request limit. Waiting {wait_time:.2f} seconds." ) time.sleep(wait_time) # Reset the counter and start time after waiting, just in case. self.request_count_minute = 0 self.minute_start_time = time.time() # Handle delay parameter time_since_last_request = current_time - self.last_request_time if self.debug_mode: self.logger.info( f"Time since last request: {time_since_last_request} seconds" ) if time_since_last_request < self.request_delay: sleep_time = self.request_delay - time_since_last_request if self.debug_mode: self.logger.info(f"Waiting {sleep_time} seconds before next request") time.sleep(sleep_time) self.last_request_time = time.time() if self.requests_per_minute > 0: self.request_count_minute += 1 def ocr(self, img_base64: str, prompt_override: str = None) -> str: """ Performs OCR on a base64 encoded image. """ if self.debug_mode: self.logger.debug(f"Starting OCR on image") self._respect_delay() try: prompt_text = prompt_override if prompt_override else self.prompt if self.language: prompt_text += f" The language is {self.language}." messages = [] if self.system_prompt: messages.append({"role": "system", "content": self.system_prompt}) messages.append( { "role": "user", "content": [ {"type": "text", "text": prompt_text}, { "type": "image_url", "image_url": { "url": f"data:image/jpeg;base64,{img_base64}" }, }, ], } ) # Determine the model model_name = self.override_model if not model_name: # If override_model is empty model_name = self.model if not model_name: # If model is also empty, determine from provider provider = self.provider # You might want to set default models for each provider here if provider == "OpenAI": model_name = "gpt-4-vision-preview" elif provider == "Google": model_name = "gemini-1.5-pro-latest" else: model_name = "gpt-4-vision-preview" # Default # Extract model name without provider prefix if it exists if ": " in model_name: model_name = model_name.split(": ", 1)[1] # Log the model being used if self.debug_mode: self.logger.info(f"Using model: {model_name}") response = self.client.chat.completions.create( model=model_name, messages=messages, max_tokens=300, # Adjust as needed ) if response.choices: full_text = response.choices[0].message.content if full_text is None: # Добавлена проверка на None if self.debug_mode: self.logger.warning("OCR response content is None.") return "" # Возвращаем пустую строку в случае None if self.debug_mode: self.logger.debug(f"OCR result: {full_text}") return full_text else: if self.debug_mode: self.logger.warning("No text found in OCR response") return "" except Exception as e: self.logger.error(f"OCR error: {e}") return "" def _ocr_blk_list( self, img: np.ndarray, blk_list: List[TextBlock], *args, **kwargs ): """ Processes a list of text blocks in an image. """ im_h, im_w = img.shape[:2] if self.debug_mode: self.logger.debug(f"Image dimensions: {im_h}x{im_w}") for blk in blk_list: x1, y1, x2, y2 = blk.xyxy if self.debug_mode: self.logger.debug(f"Processing block: ({x1}, {y1}, {x2}, {y2})") if ( y2 <= im_h and x2 <= im_w and x1 >= 0 and y1 >= 0 and x1 < x2 and y1 < y2 ): cropped_img = img[y1:y2, x1:x2] # Encode the cropped image to base64 _, buffer = cv2.imencode(".jpg", cropped_img) img_base64 = base64.b64encode(buffer).decode("utf-8") if self.debug_mode: self.logger.debug(f"Cropped image dimensions: {cropped_img.shape}") blk.text = self.ocr( img_base64, prompt_override=kwargs.get("prompt", "") ) else: if self.debug_mode: self.logger.warning("Invalid text block coordinates") blk.text = "" def ocr_img(self, img: np.ndarray, prompt: str = "") -> str: """ Performs OCR on the entire image. """ # Encode the entire image to base64 _, buffer = cv2.imencode(".jpg", img) img_base64 = base64.b64encode(buffer).decode("utf-8") return self.ocr(img_base64, prompt_override=prompt) def updateParam(self, param_key: str, param_content): super().updateParam(param_key, param_content) if param_key in [ "api_key", "endpoint", "proxy", "provider", "model", "override_model", ]: self._initialize_client() if param_key in ["requests_per_minute", "delay"]: current_time = time.time() self.request_count_minute = 0 self.minute_start_time = current_time self.last_request_time = current_time