import streamlit as st
import cv2
import numpy as np

face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

# Function to segment face from image
def segment_face(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
    for (x, y, w, h) in faces:
        face = image[y:y+h, x:x+w]
        face = cv2.resize(face, (256, 256))
        face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
        return face

def main():
    st.title("Face Segmentation")
    uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        image = np.array(bytearray(uploaded_file.read()), dtype=np.uint8)
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)

        face = segment_face(image)

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

        st.subheader("Segmented Face")
        st.image(face)

        st.subheader("Download Segmented Face")
        st.download_button(
            label="Download",
            data=cv2.imencode('.png', face)[1].tostring(),
            file_name='segmented_face.png',
            mime='image/png'
        )

if __name__ == '__main__':
    main()