import streamlit as st
import cv2
import numpy as np
from PIL import Image

def show_img(img, title="Image"):
    # Check if the image has one channel (grayscale)
    if len(img.shape) == 2:
        # Convert the grayscale image to BGR format by stacking the channels
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    st.image(img, caption=title, channels="BGR", use_column_width=True)

def find_contours(edged_img):
    contours, _ = cv2.findContours(edged_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours = sorted(contours, key=cv2.contourArea, reverse=True)
    return contours

def sort_points(pts):
    pts = pts.reshape((4, 2))
    new_pts = np.zeros((4, 2), dtype="float32")

    sum_pts = pts.sum(axis=1)
    diff_pts = np.diff(pts, axis=1)

    new_pts[0] = pts[np.argmin(sum_pts)]
    new_pts[2] = pts[np.argmax(sum_pts)]
    new_pts[1] = pts[np.argmin(diff_pts)]
    new_pts[3] = pts[np.argmax(diff_pts)]

    return new_pts

def transform_image(image_file):
    img = cv2.imread(image_file)
    original = img.copy()
    show_img(img, "Original Image")
    (H, W) = img.shape[:2]

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (7, 7), 0)
    edged = cv2.Canny(blur, 60, 160)
    show_img(edged, "Edged Image")
    conts = find_contours(edged.copy())
    for c in conts:
        peri = cv2.arcLength(c, True)
        aprox = cv2.approxPolyDP(c, 0.02 * peri, True)

        if len(aprox) == 4:
            larger = aprox
            break

    cv2.drawContours(img, [larger], -1, (120, 255, 0), 28)
    cv2.drawContours(img, [larger], -1, (120, 255, 0), 2)
    show_img(img, "Contours Image")

    points_larger = sort_points(larger)
    pts1 = np.float32(points_larger)
    pts2 = np.float32([[0, 0], [W, 0], [W, H], [0, H]])

    matrix = cv2.getPerspectiveTransform(pts1, pts2)
    transform = cv2.warpPerspective(original, matrix, (W, H))

    show_img(transform, "Transformed Image")
    return transform

def process_img(img):
    processed_img = cv2.resize(img, None, fx=1.6, fy=1.6, interpolation=cv2.INTER_CUBIC)
    processed_img = cv2.cvtColor(processed_img, cv2.COLOR_BGR2GRAY)
    processed_img = cv2.adaptiveThreshold(processed_img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 9)
    return processed_img

# Streamlit App
st.set_page_config(page_title="Document Scanner", page_icon="📄")

st.title("Document Scanner with Computer Vision")
st.write("Upload an image of a document, and the app will scan and enhance it for you.")

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = np.array(Image.open(uploaded_file))
    st.write("## Original Image")
    show_img(image)

    # Save the image temporarily for OpenCV processing
    temp_image_path = "temp_image.jpg"
    Image.fromarray(image).save(temp_image_path)

    # Transform and process the image
    transformed_image = transform_image(temp_image_path)
    processed_image = process_img(transformed_image)
    
    st.write("## Processed Image")
    show_img(processed_image)