|
import os |
|
from transformers import AutoModel, AutoTokenizer |
|
from transformers_modules.RufusRubin777.GOT_OCR2_0_CPU.modeling_GOT import GOTModel, GOTConfig |
|
|
|
class OCRModel: |
|
_instance = None |
|
|
|
def __new__(cls): |
|
if cls._instance is None: |
|
cls._instance = super(OCRModel, cls).__new__(cls) |
|
cls._instance.initialize() |
|
return cls._instance |
|
|
|
def initialize(self): |
|
model_path = os.getenv('MODEL_PATH', 'RufusRubin777/GOT-OCR2_0_CPU') |
|
|
|
|
|
config = GOTConfig.from_pretrained(model_path) |
|
self.model = GOTModel.from_pretrained( |
|
model_path, |
|
config=config, |
|
local_files_only=False |
|
) |
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained( |
|
model_path, |
|
local_files_only=False |
|
) |
|
|
|
self.model.eval() |
|
|
|
def process_image(self, image_stream): |
|
try: |
|
|
|
image = Image.open(image_stream) |
|
|
|
with torch.no_grad(): |
|
result = self.model.chat(self.tokenizer, image, ocr_type='format') |
|
return result |
|
except Exception as e: |
|
return f"Error processing image: {str(e)}" |
