Update extract_text.py

extract_text.py

@@ -2,21 +2,27 @@ 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(upload_file, gpu_available):
-    upload_file.file.seek(0)
-    file_bytes = np.frombuffer(upload_file.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)
+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))
@@ -30,4 +36,4 @@ def extract_text_from_image(upload_file, gpu_available):
         if score > 0.1:
             all_text_found.append(text)
 
-    return all_text_found
+    return all_text_found