import cv2
import streamlit as st
import tempfile
import torch
from torchvision import transforms
from mtcnn import MTCNN
from skimage.feature import hog
import joblib
import numpy as np

# Preprocessing for Siamese Model
def preprocess_image_siamese(img):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ])
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return transform(img)

# Preprocessing for SVM model (converting to grayscale)
def preprocess_image_svm(img):
    img = cv2.resize(img, (224, 224))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return img

# Extract HOG Features
def extract_hog_features(img):
    hog_features = hog(img, orientations=9, pixels_per_cell=(16, 16), cells_per_block=(4, 4))
    return hog_features

# Detect faces using MTCNN
def get_face(img):
    detector = MTCNN()
    faces = detector.detect_faces(img)
    if faces:
        x1, y1, w, h = faces[0]['box']
        x1, y1 = abs(x1), abs(y1)
        x2, y2 = x1 + w, y1 + h
        return img[y1:y2, x1:x2]
    return None

# Function to verify face (either HOG-SVM or Siamese model)
def verify(image, model, person):
    
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_image:
        temp_image.write(image.read())
        temp_image_path = temp_image.name
        
    image = cv2.imread(temp_image_path)

    face = get_face(image)

    if face is not None:
        if model_type == "HOG-SVM":
            with open(f'./svm_{lower(person)}.pkl', 'rb') as f:
                svm = joblib.load(f)
            with open(f'./pca_{lower(person)}.pkl', 'rb') as f:
                pca = joblib.load(f)

            face = preprocess_image_svm(face)

            hog = extract_hog_features(face)

            hog_pca = pca.transform([hog])

            pred = svm.predict(hog_pca)

            if pred == 1:
                st.write("Match")
            else:
                st.write("Not Match")
    else:
        st.write("Face not detected in one or both images")

# Main function to handle Streamlit interaction
def main():
    st.title("Real-time Face Verification App")
    
    model = st.selectbox("Select Model", ["Siamese", "HOG-SVM"])
    person = st.selectbox("Select Person", ["Theo"])
    enable = st.checkbox("Enable camera")
    captured_image = st.camera_input("Take a picture", disabled=not enable)

    if captured_image:
        verify(captured_image, model, person)

if __name__ == "__main__":
    main()