app.py

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

class VGGFaceEmbedding(nn.Module):
    def __init__(self):
        super(VGGFaceEmbedding, self).__init__()
        self.base_model = resnet50(pretrained=True)
        self.base_model = nn.Sequential(*list(self.base_model.children())[:-2])
        self.pooling = nn.AdaptiveAvgPool2d((1, 1))
        self.flatten = nn.Flatten()

    def forward(self, x):
        x = self.base_model(x)
        x = self.pooling(x)
        x = self.flatten(x)
        return x

class L1Dist(nn.Module):
    def __init__(self):
        super(L1Dist, self).__init__()

    def forward(self, input_embedding, validation_embedding):
        return torch.abs(input_embedding - validation_embedding)

class SiameseNetwork(nn.Module):
    def __init__(self):
        super(SiameseNetwork, self).__init__()
        self.embedding = VGGFaceEmbedding()
        self.distance = L1Dist()
        self.fc1 = nn.Linear(2048, 512)
        self.fc2 = nn.Linear(512, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, input_image, validation_image):
        input_embedding = self.embedding(input_image)
        validation_embedding = self.embedding(validation_image)
        distances = self.distance(input_embedding, validation_embedding)
        x = self.fc1(distances)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return x

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

def preprocess_image_svm(img):
    img = cv2.resize(img, (224, 224))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return img

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

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

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 == "Siamese":
            siamese = SiameseNetwork()
            siamese.load_state_dict(torch.load(f'siamese_{person.lower()}.pth'))

            face = preprocess_image_siamese(face)

            # Move to device
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            model.to(device)
            face = face.to(device)
        
            with torch.no_grad():
                output = model(face)
                probability = output.item()
                pred = 1.0 if probability > 0.5 else 0.0
        
            if pred == 1:
                st.write("Match")
            else:
                st.write("Not Match")
        
        elif model == "HOG-SVM":
            with open(f'./svm_{person.lower()}.pkl', 'rb') as f:
                svm = joblib.load(f)
            with open(f'./pca_{person.lower()}.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")

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()