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

def artistic_pencil_effect(img, sigma_s=60, sigma_r=0.07, shade_factor=0.05):
    # Convert to grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Apply a median blur to the grayscale image
    gray = cv2.medianBlur(gray, 5)

    # Use adaptive thresholding to create an edge mask
    edges = cv2.adaptiveThreshold(gray, 255,
                                  cv2.ADAPTIVE_THRESH_MEAN_C,
                                  cv2.THRESH_BINARY,
                                  blockSize=9,
                                  C=2)

    # Use cv2.stylization to create a cartoon-like effect
    stylized = cv2.stylization(img, sigma_s=sigma_s, sigma_r=sigma_r)

    # Use the edge mask to combine with the stylized image
    pencil_effect = cv2.bitwise_and(stylized, stylized, mask=edges)

    return pencil_effect

# Streamlit app layout
st.title("Artistic Pencil Effect Creator")
st.subheader("Upload an image and adjust parameters")

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

if uploaded_file is not None:
    # Read and display original image
    img = Image.open(uploaded_file)
    img = np.array(img)

    # Create a container for sliders on the left side
    with st.sidebar:
        st.header("Adjust Parameters")
        sigma_s = st.slider("Sigma S (spatial sigma)", min_value=10, max_value=200, value=60)
        sigma_r = st.slider("Sigma R (range sigma)", min_value=0.01, max_value=0.1, value=0.07, step=0.01)
        shade_factor = st.slider("Shade Factor", min_value=0.01, max_value=0.1, value=0.05, step=0.01)

    # Show a spinner while processing the pencil effect
    with st.spinner('Processing...'):
        # Apply the artistic pencil effect with selected parameters in real-time
        pencil_image = artistic_pencil_effect(img, sigma_s, sigma_r, shade_factor)

    # Display original and pencil effect images side by side on the right side
    col1, col2 = st.columns(2)
    
    with col1:
        st.image(img.astype(np.uint8), caption='Original Image', use_column_width=True)

    with col2:
        st.image(pencil_image.astype(np.uint8), caption='Artistic Pencil Effect', use_column_width=True)

    # Download button for the pencil effect image
    if st.button("Download Pencil Effect Image"):
        # Convert the pencil effect image to PIL format for download
        pencil_image_pil = Image.fromarray(pencil_image)
        pencil_image_pil.save("artistic_pencil_effect_image.png")
        
        with open("artistic_pencil_effect_image.png", "rb") as f:
            st.download_button(
                label="Download Image",
                data=f,
                file_name="artistic_pencil_effect_image.png",
                mime="image/png"
            )