app.py

import gradio as gr
from gradio_webrtc import WebRTC
import os
import pickle
import numpy as np
import cv2
import tensorflow as tf
from sklearn.preprocessing import OneHotEncoder
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras import Model
from tensorflow.image import resize
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from keras_vggface.vggface import VGGFace
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import RandomFlip, RandomRotation, RandomBrightness, RandomContrast, RandomTranslation, Input, Dense, Flatten
from sklearn.model_selection import train_test_split

class ResNetModel():
    def __init__(self):
        self.data_augmentation = Sequential([
            RandomBrightness(factor=0.2),
            RandomContrast(factor=0.2),
            RandomFlip("horizontal"),
            RandomTranslation(height_factor=0.1, width_factor=0.1),
            RandomRotation(factor=0.1)
        ])
        
    def load_data(self, data_path: str, input_size:tuple = (224, 224)):
        image_list = []
        class_list = []
        self.input_size = input_size
        faceCascade = cv2.CascadeClassifier('Cascades/haarcascade_frontalface_default.xml')
        for label in os.listdir(data_path):
            if len(os.listdir(os.path.join(data_path, label))) < 5:
                continue
            for j, filename in enumerate(os.listdir(os.path.join(data_path, label))):
                if j >= 10:
                    break
                filename = os.path.join(data_path, label, filename)
                print(f"Index {j}. {filename}")
                image = load_img(filename)
                # image = load_img(filename, color_mode = color_mode)
                # if color_mode != "grayscale":
                #     gray = np.array(rgb_to_grayscale(image))
                # else:
                #     gray = np.array(image.copy())
                image = np.array(image)
                faces = faceCascade.detectMultiScale(
                    image,     
                    scaleFactor=1.2,
                    minNeighbors=5,     
                    minSize=(20, 20)
                )
                for (x,y,w,h) in faces:
                    if w == 0 or h == 0:
                        continue
                    image_roi = image[y:y+h, x:x+w]
                    image_roi = img_to_array(image_roi)
                    image_roi = resize(image_roi, input_size)
                    
                    image_list.append(image_roi)
                    class_list.append(label)
                    
        os.system("cls")
        encoder = OneHotEncoder(sparse_output=False)
        class_list = encoder.fit_transform(np.array(class_list).reshape(-1, 1))
        
        image_list = np.asarray(image_list)
        self.label_names = encoder.categories_[0]
        
        X_train, X_temp, y_train, y_temp = train_test_split(image_list, class_list, test_size=0.3)
        X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.5)
        
        self.X_train = X_train
        self.y_train = y_train
        self.X_val = X_val
        self.y_val = y_val
        self.X_test = X_test
        self.y_test = y_test
        print(f"Train dataset len: {X_train.shape[0]}")
        print(f"Val dataset len: {X_val.shape[0]}")
        print(f"Test dataset len: {X_test.shape[0]}")
        print(f"Sample image shape: {X_train[0].shape}")
                    
    def save(self, save_path):
        self.X_train = []
        self.y_train = []
        self.X_val = []
        self.y_val = []
        self.X_test = []
        self.y_test = []
        
        model_save_path = save_path + "_model.h5"
        self.model.save(model_save_path)

        # Save the rest of the class attributes
        with open(save_path + "_attributes.pickle", "wb") as file:
            pickle.dump({
                "data_augmentation": self.data_augmentation,
                "label_names": self.label_names,
                "input_size": self.input_size,
            }, file)
            
    @classmethod
    def load(cls, save_path):
        with open(save_path + "_attributes.pickle", "rb") as file:
            attributes = pickle.load(file)

        instance = cls()
        instance.data_augmentation = attributes["data_augmentation"]
        instance.label_names = attributes["label_names"]
        instance.input_size = attributes["input_size"]

        # Load the Keras model
        model_save_path = save_path + "_model.h5"
        instance.model = tf.keras.models.load_model(model_save_path)

        print(f"Model and attributes loaded from {save_path}")
        return instance

        
    def fit(self, model_name: str, augmentation=True, save_path=None, batch_size=64, epochs=10):
        inputs = Input(shape=self.X_train[0].shape)
        layer = self.data_augmentation(inputs) if augmentation else inputs
        
        base_model = VGGFace(model=model_name, include_top=False, input_shape=self.X_train[0].shape, pooling="avg")
        base_model.trainable = False
        
        layer = base_model(layer)
        
        layer = Flatten(name="Flatten")(layer)
        out = Dense(len(self.label_names), activation="softmax")(layer)
        
        model = Model(inputs, out)
        model.compile(optimizer="adam", metrics=["accuracy"], loss="categorical_crossentropy")
        
        print(model.summary())
        
        self.history = model.fit(
            self.X_train, self.y_train, 
            batch_size=batch_size, 
            epochs=epochs, 
            validation_data=(self.X_val, self.y_val),
            callbacks=[EarlyStopping(monitor="val_loss", patience=3)]
        )
        
        self.model = model
        if save_path is not None:
            self.save(save_path)
        
        loss, accuracy = model.evaluate(self.X_test, self.y_test)
        print(f"Test accuracy:{(accuracy * 100):2f}")
        print(f"Test loss:{loss:2f}")
        return model

            
    def predict(self, predict_image):
        predict_image = img_to_array(predict_image)
        predict_image = resize(predict_image, self.input_size)
        predict_image = np.expand_dims(predict_image, axis=0)
        predict_label = self.model.predict(predict_image)
        return self.label_names[np.argmax(predict_label)], np.max(predict_label)

faceRecognition = ResNetModel.load("resnet_v2")
def predict(image):
    faceCascade = cv2.CascadeClassifier('Cascades/haarcascade_frontalface_default.xml')
    frame = np.array(image)
    faces = faceCascade.detectMultiScale(
        frame,     
        scaleFactor=1.2,
        minNeighbors=5,     
        minSize=(20, 20)
    )
    for (x,y,w,h) in faces:
        cv2.rectangle(frame,(x,y),(x+w,y+h),(255,0,0),3)
        roi_color = frame[y:y+h, x:x+w]
        roi_color = cv2.flip(roi_color, 1)
        label, conf = faceRecognition.predict(roi_color)
        label = label[:12] + "..." if len(label) > 10 else label
        text = "{label} : {conf:.3f}".format(label = label, conf = conf)
        
        cv2.putText(frame, text, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,0,0), 3)
    return frame

css = """.my-column {max-width: 600px;}"""
with gr.Blocks(css=css) as demo:
    gr.HTML(
        """
        <h1 style='text-align: center'>
            Face Recognition Using ResNet-50 LB01 - Kelompok 8
        </h1>
        <ul style='text-align: center'>
            <h4> <li>2602082452 - Rendy Susanto </li> </h4>
            <h4> <li>2602090031 - Rafael Juviano Joesoef </li> </h4> 
            <h4> <li>2602091122 - Owen Tamashi Buntoro </li> </h4>
        </ul>
        """
    )
    with gr.Row():
        with gr.Column(elem_classes=["my-column"]):
            input_img = gr.Image(label="Input", sources="webcam")
        with gr.Column(elem_classes=["my-column"]):
            output_img = gr.Image(label="Output")

        input_img.stream(
            fn=predict, inputs=[input_img], outputs=[output_img], time_limit=1, stream_every=0.1
        )

if __name__ == "__main__":
    demo.launch()