extract_text.py

import numpy as np
import cv2
import easyocr
import torch
from starlette.datastructures import UploadFile

device = "cuda" if torch.cuda.is_available() else "cpu"
reader = easyocr.Reader(["en"], gpu=(device == "cuda"), verbose=False)

def extract_text_from_image(image_input, gpu_available):
    # Se for UploadFile, lê bytes
    if isinstance(image_input, UploadFile):
        image_input.file.seek(0)
        file_bytes = np.frombuffer(image_input.file.read(), np.uint8)
        img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError("Não foi possível decodificar a imagem.")
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    # Se já for ndarray
    elif isinstance(image_input, np.ndarray):
        img = image_input
    else:
        raise ValueError("Formato de imagem não suportado.")

    # Resto do processamento
    scale_factor = 2
    upscaled = cv2.resize(img, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR)
    blur_img = cv2.blur(upscaled, (5, 5))
    
    reader = easyocr.Reader(['en'], gpu=gpu_available, verbose=False)
    all_text_found = []
    text_ = reader.readtext(blur_img, detail=1, paragraph=False, text_threshold=0.3)

    for t in text_:
        bbox, text, score = t
        if score > 0.1:
            all_text_found.append(text)

    return all_text_found