File size: 2,636 Bytes
51c49bc 26f855a 2795ce6 237de6f 6139662 51c49bc 26f855a 6139662 c083a98 d20f881 26f855a 6101096 c083a98 6139662 c083a98 6101096 26f855a c083a98 2795ce6 c083a98 2795ce6 26f855a 6139662 51c49bc |
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 |
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image
import cv2
import os
import torch
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils import notification_queue, log_print
from all_models import models
def text(image_cv):
try:
# Get model instance from singleton
model, processor = models.get_trocr_model()
if not isinstance(image_cv, list):
image_cv = [image_cv]
t = ""
total_images = len(image_cv)
log_print(f"Processing {total_images} image(s) for text extraction")
for i, img in enumerate(image_cv):
try:
log_print(f"Processing image {i+1}/{total_images}")
# Validate image
if img is None:
log_print(f"Skipping image {i+1} - Image is None", "WARNING")
continue
# Convert to RGB
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image = Image.fromarray(img_rgb)
# Get pixel values
pixel_values = processor(image, return_tensors="pt").pixel_values
if torch.cuda.is_available():
pixel_values = pixel_values.to(models.device)
# Generate text
generated_ids = model.generate(pixel_values)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
# Clean up the text
cleaned_text = generated_text.replace(" ", "")
t = t + cleaned_text + " "
log_print(f"Successfully extracted text from image {i+1}: {cleaned_text}")
# Clean up CUDA memory
if torch.cuda.is_available():
del pixel_values
del generated_ids
torch.cuda.empty_cache()
except Exception as e:
log_print(f"Error processing image {i+1}: {str(e)}", "ERROR")
continue
return t.strip()
except Exception as e:
error_msg = f"Error in text function: {str(e)}"
log_print(error_msg, "ERROR")
notification_queue.put({
"type": "error",
"message": error_msg
})
return ""
finally:
# Release model reference
models.release_trocr_model()
|