import streamlit as st
import cv2
import numpy as np
import tempfile
from PIL import Image
from ultralytics import YOLO

def process_lines(image_path):
    thickness = 3

    image = cv2.imread(image_path)
    result = image.copy()

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=80, minLineLength=40, maxLineGap=40)
    line_mask = np.zeros_like(gray)

    if lines is not None:
        for line in lines:
            x1, y1, x2, y2 = line[0]
            cv2.line(line_mask, (x1, y1), (x2, y2), (255, 255, 255), thickness=3)

    return line_mask

def detect_text(image_path):
    model = YOLO("best 3.pt")
    results = model.predict(image_path)
    annotated_image = results[0].plot()
    return annotated_image

st.title("Line and Text Extraction")
st.sidebar.header("Upload an Image")

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

if st.sidebar.button("Process Image"):
    if uploaded_file is not None:
        with tempfile.NamedTemporaryFile(delete=False, suffix=uploaded_file.name) as temp_file:
            temp_file.write(uploaded_file.read())
            temp_file_path = temp_file.name

        line_mask = process_lines(temp_file_path)
        text_extracted=detect_text(temp_file_path)

        st.subheader("Original image")
        st.image(uploaded_file)

        st.subheader("Lines Extracted")
        st.image(line_mask, channels="GRAY")
        
        st.subheader("Text Detected")
        st.image(cv2.cvtColor(text_extracted, cv2.COLOR_BGR2RGB))
    else:
        st.sidebar.error("Please upload an image file before clicking 'Process Image'.")