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
from transformers import pipeline  # Importing the pipeline

"""
Paddle OCR
"""
def ocr_with_paddle(img):
    finaltext = ''
    ocr = PaddleOCR(lang='en', use_angle_cls=True)
    result = ocr.ocr(img)
    for i in range(len(result[0])):
        text = result[0][i][1][0]
        finaltext += ' ' + 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)
    first = predictions[0]
    for text, box in first:
        output_text += ' ' + text
    return output_text

"""
easy OCR
"""
# grayscale 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(['th', 'en'])
    bounds = reader.readtext('image.png', paragraph="False", detail=0)
    bounds = ''.join(bounds)
    return bounds

"""
Generate OCR
"""
def generate_ocr(Method, img):
    text_output = ''
    if (img).any():
        print("Method___________________", Method)
        if Method == 'EasyOCR':
            text_output = ocr_with_easy(img)
        if Method == 'KerasOCR':
            text_output = ocr_with_keras(img)
        if Method == 'PaddleOCR':
            text_output = ocr_with_paddle(img)
        
        try:
            flag(Method, text_output, img)
        except Exception as e:
            print(e)
        
        # Generate Text using FLAN-T5 model
        text_gen = generate_text_with_flan_t5(text_output)
        return text_gen
    else:
        raise gr.Error("Please upload an image!!!!")

"""
Text Generation using FLAN-T5
"""
def generate_text_with_flan_t5(input_text):
    # Load the pre-trained FLAN-T5 model
    pipe = pipeline("text2text-generation", model="google/flan-t5-large")

    # Use the model to generate a response based on the OCR output
    output = pipe(input_text)
    return output[0]['generated_text']


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

image = gr.Image()
method = gr.Radio(["PaddleOCR", "EasyOCR", "KerasOCR"], value="PaddleOCR")
output = gr.Textbox(label="Generated Text")

demo = gr.Interface(
    generate_ocr,
    [method, image],
    output,
    title="Optical Character Recognition and Text Generation",
    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()