import cv2
import easyocr
import gradio as gr
import base64
import json

def text_extraction(image):
    # Convert base64 image to OpenCV format
    image = base64.b64decode(image.split(",")[1])
    nparr = np.frombuffer(image, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    # Instance text detector
    reader = easyocr.Reader(['en'], gpu=False)

    # Detect text on image
    text_ = reader.readtext(img)

    threshold = 0.25
    # Draw bbox and text
    for t_, t in enumerate(text_):
        bbox, text, score = t

        if score > threshold:
            cv2.rectangle(img, tuple(map(int, bbox[0])), tuple(map(int, bbox[2])), (255, 0, 0), 2)

    # Encode image to base64
    retval, buffer = cv2.imencode('.jpg', img)
    img_base64 = base64.b64encode(buffer).decode('utf-8')

    # Create JSON response
    response_json = {
        'annotated_image_base64': img_base64,
        'text_data': text_
    }

    # Convert the dictionary to a JSON string
    response_json_str = json.dumps(response_json, default=str)

    return response_json_str

# Define Gradio interface
iface = gr.Interface(
    fn=text_extraction,
    inputs=gr.Image(),
    outputs=["image", "json"]
)

# Launch the Gradio interface
iface.launch()