import streamlit as st
import cv2
from PIL import Image,ImageColor
import io
import numpy as np
from passport_img import merge_all,predict_on_image,overlay,add_name_dob,align_crop_image


def grayscale_image(image):
    # Convert the image to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    return gray
st.title('📷Passport Photo Maker📷')
st.title('🤖AI Enabled🤖')
st.write("Made with ❤️ by Mainak")
def main():
    # st.title("Grayscale Image Converter")

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

    if uploaded_file is not None:
        addname = st.checkbox('Select if you want to add name and Date of Birth')
        aicrop = st.checkbox('Select if you want to crop or change background using AI')
        size = st.selectbox('Select Size',('3X4','3.5X4.5'))
        size = int(''.join(i for i in size if not i in ['X','.']))
        if size==3545:
            num_img = st.slider('Number of Image', 0, 36, 6)
        elif size==34:
            num_img = st.slider('Number of Image', 0, 49, 7)
        if addname and not aicrop:
            col1,col2 = st.columns(2)
            with col1:  
                # st.text('Enter Name')
                name = st.text_input('Enter Name',value="First_name Last_name",key = 1)
            with col2:
                # st.text('Enter Date of Birth')
                dob = st.text_input('Enter Date of Birth',value="01/01/0001",key = 2)
           
            if st.button('make'):
                with st.spinner('Ruko jara sabar karo....'):
                    image = Image.open(uploaded_file)
                    img_array = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
                    img_array = add_name_dob(img_array,name,dob)
                    gray_image = merge_all(num_img,img_array,size)
                    gray_image = cv2.cvtColor(np.array(gray_image), cv2.COLOR_BGR2RGB)

                    pil_image = Image.fromarray(gray_image)
                    st.image(pil_image, caption='Passport Image', use_column_width=True)

                    # Create a download link for the grayscale image
                    buffered = io.BytesIO()
                    pil_image.save(buffered, format="PNG")
                    download_btn = st.download_button(
                        label='Download Image',
                        data=buffered.getvalue(),
                        file_name='passport_image.png',
                        mime='image/png'
                    )
        elif addname and aicrop:
            color = st.color_picker('Pick A Color for Background', '#00f900')
            color = ImageColor.getcolor(color, "RGB")

            col1,col2 = st.columns(2)
            with col1:  
                # st.text('Enter Name')
                name = st.text_input('Enter Name',value="First_name Last_name",key = 1)
            with col2:
                # st.text('Enter Date of Birth')
                dob = st.text_input('Enter DOB',value="01/01/0001",key = 2)

            if st.button('make'):
                with st.spinner('Ruko jara sabar karo....'):
                    image = Image.open(uploaded_file)
                    img_array = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
                    img = align_crop_image(img_array)
                    if img is not None:
                        # predict by YOLOv8
                        if predict_on_image(img, conf=0.6) is not None:
                            boxes, masks, cls, probs = predict_on_image(img, conf=0.6)
                            # overlay masks on original image
                            image_with_masks = np.copy(img)
                            for mask_i in masks:
                                image_with_masks = overlay(image_with_masks, mask_i, color=color, alpha=1)
                        else:
                            image_with_masks=img
                        img_array = add_name_dob(image_with_masks,name,dob)
                        gray_image = merge_all(num_img,img_array,size)
                        gray_image = cv2.cvtColor(np.array(gray_image), cv2.COLOR_BGR2RGB)

                        pil_image = Image.fromarray(gray_image)
                        st.image(pil_image, caption='Passport Image', use_column_width=True)

                        # Create a download link for the grayscale image
                        buffered = io.BytesIO()
                        pil_image.save(buffered, format="PNG")
                        download_btn = st.download_button(
                            label='Download Image',
                            data=buffered.getvalue(),
                            file_name='passport_image.png',
                            mime='image/png'
                        )
                    else:
                        st.warning('Try with another image')
        elif aicrop and not addname:
            
            color = st.color_picker('Pick A Color for Background', '#00f900')
            color = ImageColor.getcolor(color, "RGB")
            if st.button('make'):
                with st.spinner('Ruko jara sabar karo....'):
                    image = Image.open(uploaded_file)
                    img_array = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
                    img = align_crop_image(img_array)
                    if img is not None:
                        # predict by YOLOv8
                        if predict_on_image(img, conf=0.6) is not None:
                            boxes, masks, cls, probs = predict_on_image(img, conf=0.6)
                            # overlay masks on original image
                            image_with_masks = np.copy(img)
                            for mask_i in masks:
                                image_with_masks = overlay(image_with_masks, mask_i, color=color, alpha=1)
                        else:
                            image_with_masks=img
                        # img_array = add_name_dob(image_with_masks,name,dob)
                        gray_image = merge_all(num_img,image_with_masks,size)
                        gray_image = cv2.cvtColor(np.array(gray_image), cv2.COLOR_BGR2RGB)

                        pil_image = Image.fromarray(gray_image)
                        st.image(pil_image, caption='Passport Image', use_column_width=True)

                        # Create a download link for the grayscale image
                        buffered = io.BytesIO()
                        pil_image.save(buffered, format="PNG")
                        download_btn = st.download_button(
                            label='Download Image',
                            data=buffered.getvalue(),
                            file_name='passport_image.png',
                            mime='image/png'
                        )
                    else:
                        st.error('Try with another image')

        else:
            if st.button('make'):
                with st.spinner('Ruko jara sabar karo....'):
                    image = Image.open(uploaded_file)
                    img_array = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
                    gray_image = merge_all(num_img,img_array,size)
                    gray_image = cv2.cvtColor(np.array(gray_image), cv2.COLOR_BGR2RGB)

                    pil_image = Image.fromarray(gray_image)
                    st.image(pil_image, caption='Passport Image', use_column_width=True)

                    # Create a download link for the grayscale image
                    buffered = io.BytesIO()
                    pil_image.save(buffered, format="PNG")
                    download_btn = st.download_button(
                        label='Download Image',
                        data=buffered.getvalue(),
                        file_name='passport_image.png',
                        mime='image/png'
                    )
        

if __name__ == '__main__':
    main()