import gradio as gr
import requests
import os
from datasets import load_dataset, Image
from PIL import Image
from paddleocr import PaddleOCR
from doctr.io import DocumentFile
import torch

# Set environment variable for PyTorch usage
os.environ['USE_TF'] = '0'  # Set TensorFlow to off
os.environ['USE_TORCH'] = '1'  # Set PyTorch to on

from doctr.models import ocr_predictor
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Example for PyTorch model or doctr model

ocr_model = ocr_predictor(det_arch='db_mobilenet_v3_large', reco_arch='crnn_mobilenet_v3_small', pretrained=True).to(device)


import torch

# Check if CUDA is available
if torch.cuda.is_available():
    print(f"GPU is available. Device: {torch.cuda.get_device_name(0)}")
else:
    print("GPU is not available, using CPU.")

    
"""
Perform OCR with doctr
"""
def ocr_with_doctr(file):
    text_output = ''
    
    # Load the document
    doc = DocumentFile.from_pdf(file)
    
    # Perform OCR
    result = ocr_model(doc)
    
    # Extract text from OCR result
    for page in result.pages:
        for block in page.blocks:
            for line in block.lines:
                text_output += " ".join([word.value for word in line.words]) + "\n"
    
    return text_output

"""
Paddle OCR
"""
def ocr_with_paddle(img):
    finaltext = ''
    ocr = PaddleOCR(lang='en', use_angle_cls=True, use_gpu=True)
    # img_path = 'exp.jpeg'
    result = ocr.ocr(img)
    
    for i in range(len(result[0])):
        text = result[0][i][1][0]
        finaltext += ' '+ text
    return finaltext

def generate_ocr(Method, file):
    text_output = ''
    if isinstance(file, bytes):  # Handle file uploaded as bytes
        file = io.BytesIO(file)

    if file.name.endswith('.pdf'):
        # Perform OCR on the PDF using doctr
        text_output = ocr_with_doctr(file)

    else:
        # Handle image file
        img_np = np.array(Image.open(file))
        text_output = generate_text_from_image(Method, img_np)
    
    return text_output

def generate_text_from_image(Method, img):
    text_output = ''
    if Method == 'PaddleOCR':
        text_output = ocr_with_paddle(img)
    return text_output


import gradio as gr

image_or_pdf = gr.File(label="Upload an image or PDF")
method = gr.Radio(["PaddleOCR"], value="PaddleOCR")
output = gr.Textbox(label="Output")

demo = gr.Interface(
    generate_ocr,
    [method, image_or_pdf],
    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(share=True)
# import os

# # Disable TensorFlow to ensure PyTorch is used
# os.environ['USE_TF'] = '0'

# import torch
# print(torch.cuda.is_available())  # Should return True if GPU is available