import streamlit as st
import matplotlib.pyplot as plt
import cv2
import easyocr
import numpy as np
from PIL import Image as PILImage
from io import BytesIO
import pdf2image

# Function to handle file upload and OCR processing
def process_image(image, language):
    reader = easyocr.Reader([language])  # Initialize OCR reader with chosen language
    
    # Convert the image to OpenCV format
    open_cv_image = np.array(image)
    open_cv_image = open_cv_image[:, :, ::-1].copy()  # Convert RGB to BGR for OpenCV
    
    # Apply the scan effect
    processed_image = scan_effect(open_cv_image)
    
    # Perform OCR
    output = reader.readtext(processed_image)
    
    return processed_image, output

# Function for image scan effect
def scan_effect(img):
    blackPoint = 66
    whitePoint = 130
    image = highPassFilter(img, kSize=51)
    image_white = whitePointSelect(image, whitePoint)
    img_black = blackPointSelect(image_white, blackPoint)
    image = blackPointSelect(img, blackPoint)
    white = whitePointSelect(image, whitePoint)
    img_black = blackAndWhite(white)
    return img_black

# Helper functions for scan effect
def map(x, in_min, in_max, out_min, out_max):
    return (x - in_min) * (out_max - out_min) / (in_max - in_min) + out_min

def highPassFilter(img, kSize):
    if not kSize % 2:
        kSize += 1
    kernel = np.ones((kSize, kSize), np.float32) / (kSize * kSize)
    filtered = cv2.filter2D(img, -1, kernel)
    filtered = img.astype('float32') - filtered.astype('float32')
    filtered = filtered + 127 * np.ones(img.shape, np.uint8)
    filtered = filtered.astype('uint8')
    return filtered

def blackPointSelect(img, blackPoint):
    img = img.astype('int32')
    img = map(img, blackPoint, 255, 0, 255)
    _, img = cv2.threshold(img, 0, 255, cv2.THRESH_TOZERO)
    img = img.astype('uint8')
    return img

def whitePointSelect(img, whitePoint):
    _, img = cv2.threshold(img, whitePoint, 255, cv2.THRESH_TRUNC)
    img = img.astype('int32')
    img = map(img, 0, whitePoint, 0, 255)
    img = img.astype('uint8')
    return img

def blackAndWhite(img):
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    (l, a, b) = cv2.split(lab)
    img = cv2.add(cv2.subtract(l, b), cv2.subtract(l, a))
    return img

# Streamlit app layout
st.title("OCR and Image Processing App")

# Sidebar for file upload and language selection
st.sidebar.title("Upload and Settings")
uploaded_file = st.sidebar.file_uploader("Upload PDF or Image", type=["pdf", "png", "jpg", "jpeg"])
language_option = st.sidebar.radio("Select OCR Language", ('ne', 'en'))  # 'ne' for Nepali, 'en' for English

if uploaded_file:
    file_type = uploaded_file.type
    
    if file_type == "application/pdf":
        # Convert PDF to images
        images = pdf2image.convert_from_bytes(uploaded_file.read())
        
        st.image(images, caption='Uploaded PDF as Images', use_column_width=True)
        
        if st.button("Process PDF"):
            # Process the first page of the PDF
            st.write("Processing PDF...")
            for page_image in images:
                with st.spinner('Processing...'):
                    processed_image, ocr_output = process_image(page_image, language_option)
                    st.image(processed_image, caption="Processed Image", use_column_width=True)
                    
                    # Display OCR output
                    st.write("Extracted Text:")
                    for item in ocr_output:
                        st.write(item[1])
                        
                    # Allow download of processed image
                    img = PILImage.fromarray(processed_image)
                    buf = BytesIO()
                    img.save(buf, format="PNG")
                    byte_im = buf.getvalue()
                    st.download_button(label="Download Processed Image", data=byte_im, file_name="processed_image.png", mime="image/png")
                    
    else:
        # Handle image files
        image = PILImage.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_column_width=True)
        
        if st.button("Process Image"):
            # Process the uploaded image
            with st.spinner('Processing...'):
                processed_image, ocr_output = process_image(image, language_option)
                st.image(processed_image, caption="Processed Image", use_column_width=True)
                
                # Display OCR output
                st.write("Extracted Text:")
                for item in ocr_output:
                    st.write(item[1])
                
                # Allow download of processed image
                img = PILImage.fromarray(processed_image)
                buf = BytesIO()
                img.save(buf, format="PNG")
                byte_im = buf.getvalue()
                st.download_button(label="Download Processed Image", data=byte_im, file_name="processed_image.png", mime="image/png")