import streamlit as st
import cv2
import numpy as np
from transformers import pipeline

# Set up the CLIP classifier
model_name = "openai/clip-vit-large-patch14-336"
classifier = pipeline("zero-shot-image-classification", model=model_name)
labels_for_classification = ["genuine face", "morphed face"]

# Define the restoration function
def restore(image):
    # Convert the image to float32
    image = np.float32(image)
    # Apply a median blur to the image
    image = cv2.medianBlur(image, 5)
    return image

# Define the enhancement function
def enhance(image):
    # Convert the image to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # Apply histogram equalization to the image
    equalized = cv2.equalizeHist(gray)
    return equalized

# Define the Streamlit app
def app():
    # Create a file uploader
    uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
    if uploaded_file is not None:
        # Read the image file
        image = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)

        # Classify the image using CLIP
        scores = classifier(image, candidate_labels=labels_for_classification)
        if scores[0]['label'] == "genuine face":
            st.write("The image contains a genuine face")
        else:
            st.write("The image contains a morphed face")

        # Restore the image
        restored = restore(image)
        st.image(restored, caption="Restored Image")

        # Enhance the image
        enhanced = enhance(restored)
        st.image(enhanced, caption="Enhanced Image")

# Run the app
if __name__ == "__main__":
    app()