import gradio as gr
import tensorflow as tf
import keras_ocr
import requests
import cv2
import os
import csv
import numpy as np
import pandas as pd
import huggingface_hub
from huggingface_hub import Repository
from datetime import datetime
import scipy.ndimage.interpolation as inter
import easyocr
import datasets
from datasets import load_dataset, Image
from PIL import Image
from paddleocr import PaddleOCR
from save_data import flag
            
"""
Paddle OCR
"""
def ocr_with_paddle(img):
    finaltext = ''
    ocr = PaddleOCR(lang='en')  # Không hỗ trợ đa ngôn ngữ kiểu 'vi|en', phải chạy 2 lần
    result_en = ocr.ocr(img)

    ocr_vi = PaddleOCR(lang='vi')
    result_vi = ocr_vi.ocr(img)

    def extract_text(result):
        return ' '.join([line[1][0] for line in result[0]])

    en_text = extract_text(result_en)
    vi_text = extract_text(result_vi)

    finaltext = f"[EN]: {en_text}\n[VI]: {vi_text}"
    return finaltext

"""
Keras OCR
"""
def ocr_with_keras(img):
    output_text = ''
    pipeline = keras_ocr.pipeline.Pipeline()
    images = [keras_ocr.tools.read(img)]
    predictions = pipeline.recognize(images)
    for text, box in predictions[0]:
        output_text += ' ' + text
    return "[Detected]: " + output_text

"""
easy OCR
"""
# gray scale image
def get_grayscale(image):
    return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# Thresholding or Binarization
def thresholding(src):
    return cv2.threshold(src,127,255, cv2.THRESH_TOZERO)[1]
def ocr_with_easy(img):
    gray_scale_image = get_grayscale(img)
    thresholding(gray_scale_image)
    cv2.imwrite('image.png', gray_scale_image)
    reader = easyocr.Reader(['vi', 'en'])  # Hỗ trợ tiếng Việt và tiếng Anh
    bounds = reader.readtext('image.png', paragraph=False, detail=0)
    result_text = '\n'.join(bounds)
    return result_text
        
"""
Generate OCR
"""
def generate_ocr(Method, img):
    if img is None or not (img).any():
        raise gr.Error("Please upload an image!")

    text_output = ''
    print("Method selected:", Method)
    
    if Method == 'EasyOCR':
        text_output = ocr_with_easy(img)
    elif Method == 'KerasOCR':
        text_output = ocr_with_keras(img)
    elif Method == 'PaddleOCR':
        text_output = ocr_with_paddle(img)

    try:
        flag(Method, text_output, img)
    except Exception as e:
        print("Flag error:", e)
    
    return text_output
    
    # except Exception as e:
    #     print("Error in ocr generation ==>",e)
    #     text_output = "Something went wrong"
    # return text_output
    

"""
Create user interface for OCR demo
"""

# image = gr.Image(shape=(300, 300))
image = gr.Image()
method = gr.Radio(["PaddleOCR","EasyOCR", "KerasOCR"],value="PaddleOCR")
output = gr.Textbox(label="Output")

demo = gr.Interface(
    generate_ocr,
    [method,image],
    output,
    title="Optical Character Recognition",
    css=".gradio-container {background-color: lightgray} #radio_div {background-color: #FFD8B4; font-size: 40px;}",
    article = """<p style='text-align: center;'>Feel free to give us your thoughts on this demo and please contact us at 
                    <a href="mailto:letstalk@pragnakalp.com" target="_blank">letstalk@pragnakalp.com</a> 
                    <p style='text-align: center;'>Developed by: <a href="https://www.pragnakalp.com" target="_blank">Pragnakalp Techlabs</a></p>"""
    

)
# demo.launch(enable_queue = False)
demo.launch(show_error=True)