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

# Function to convert image to sketch with adjustable outline thickness
def image_to_sketch(image, kernel_size=1):
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    inverted_image = 255 - gray_image
    blurred_image = cv2.GaussianBlur(inverted_image, (21, 21), 0)
    inverted_blurred = 255 - blurred_image
    sketch = cv2.divide(gray_image, inverted_blurred, scale=256.0)
    
    # Apply adaptive thresholding to enhance edges
    adaptive_thresh = cv2.adaptiveThreshold(sketch, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 9, 2)
    
    # Apply morphological operation to thicken the outlines
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    sketch = cv2.dilate(adaptive_thresh, kernel, iterations=1)
    
    return sketch

# Function to blur the background of the image
def blur_background(image, blur_strength=31):
    # # Convert the image to grayscale
    # gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # # Use thresholding to create a binary mask
    # _, mask = cv2.threshold(gray, 120, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    
    # # Create the inverse mask
    # mask_inv = cv2.bitwise_not(mask)
    
    # # Blur the entire image
    # blurred = cv2.GaussianBlur(image, (blur_strength, blur_strength), 0)
    
    # # Use the mask to combine the original image with the blurred image
    # background = cv2.bitwise_and(blurred, blurred, mask=mask)
    # foreground = cv2.bitwise_and(image, image, mask=mask_inv)
    # combined = cv2.add(background, foreground)
    ksize = (10, 10) 
  
# Using cv2.blur() method  
    combined = cv2.blur(image, ksize)  
    
    return combined

# Streamlit app layout
st.set_page_config(page_title="Image to Sketch Converter", page_icon="🎨", layout="centered")

# Custom CSS for heading color and footer positioning
st.markdown("""
    <style>
    .title {
        color: blue;
        font-size: 2.5em;
        font-weight: bold;
        text-align: center;
    }
    .footer {
        position: relative;
        bottom: 0;
        width: 100%;
        background-color: orange;
        text-align: center;
        color: black;
        padding: 10px;
        font-weight: bold;
        margin-top: 50px;
    }
    .content {
        margin-bottom: 70px;
    }
    .spacing {
        margin: 10px 10px;
    }
    .centered-button {
        display: flex;
        justify-content: center;
        align-items: center;
        gap: 10px;
    }
    </style>
    """, unsafe_allow_html=True)

# Title and description
st.markdown('<p class="title">🎨 Image to Sketch Converter</p>', unsafe_allow_html=True)
st.markdown("""
    Convert your images into beautiful sketches with this simple app.
    Upload an image, and get the sketch version instantly! You can even download the sketch.
    """)

# Example conversions
st.subheader("Example Conversions")

# Load and display example image
example_image_path = 'Dog.jpg'
example_image = cv2.imread(example_image_path)

if example_image is not None:
    # Convert BGR to RGB for correct color display
    example_image_rgb = cv2.cvtColor(example_image, cv2.COLOR_BGR2RGB)
    example_image_blurred = blur_background(example_image)
    example_sketch = image_to_sketch(example_image_blurred)

    col1, col2 = st.columns(2)

    with col1:
        st.image(example_image_rgb, caption='Original Image', use_column_width=True)
    with col2:
        st.image(example_image_blurred, caption='Sketch Image', use_column_width=True)
else:
    st.error(f"Failed to load example image from path: {example_image_path}")

# User upload section
st.subheader("Upload Your Image")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Load the image
    image = Image.open(uploaded_file)
    image_np = np.array(image)
    
    # Determine the format of the uploaded image
    image_format = image.format.lower()
    
    st.write("Converting...")

    # Blur the background
    image_blurred = blur_background(image_np)
    
    # Convert the image to a sketch
    sketch = image_to_sketch(image_blurred)
    
    col3, col4 = st.columns(2)

    with col3:
        st.image(image, caption='Uploaded Image', use_column_width=True)
    with col4:
        st.image(image_blurred, caption='Sketch', use_column_width=True)

    # Add some space before the button
    st.markdown('<div class="spacing"></div>', unsafe_allow_html=True)
    
    # Convert the sketch to an image and save to an in-memory file object
    sketch_image = Image.fromarray(sketch)
    buf = io.BytesIO()
    sketch_image.save(buf, format=image_format.upper())
    byte_im = buf.getvalue()
    
    # Provide a download link for the sketch image in the center
    st.markdown('<div class="centered-button">', unsafe_allow_html=True)
    btn = st.download_button(
        label="Download Sketch",
        data=byte_im,
        file_name=f"sketch.{image_format}",
        mime=f"image/{image_format}"
    )
    st.markdown('</div>', unsafe_allow_html=True)
else:
    st.info("Please upload an image to convert.")

# Footer
st.markdown("""
    <div class="footer">
        Made by Mallela Preethi
    </div>
    """, unsafe_allow_html=True)