app.py

# Transformers and its models
import transformers

# For Image Processing
from transformers import ViTImageProcessor

# For Model
from transformers import ViTModel, ViTConfig, pipeline

# For data augmentation
from torchvision import transforms, datasets

# For GPU
from transformers import set_seed
from torch.optim import AdamW
from accelerate import Accelerator, notebook_launcher

# For Data Loaders
import datasets
from torch.utils.data import Dataset, DataLoader

# For Display
#from tqdm.notebook import tqdm

# Other Generic Libraries
import torch
from PIL import Image
import os
import streamlit as st
import gc
from glob import glob
import shutil
import pandas as pd
import numpy as np
#import matplotlib.pyplot as plt
from io import BytesIO
import torch.nn.functional as F

# Set the device (GPU or CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Initialse Globle Variables
MODEL_TRANSFORMER = 'google/vit-base-patch16-224'
BATCH_SIZE = 8

# Set Paths
data_path = 'employees'
model_path = 'vit_pytorch_GPU_1.pt'
webcam_path = 'captured_image.jpg'

# Set Title
st.title("Employee Attendance System")
#pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")

# Define Image Processor
image_processor_prod = ViTImageProcessor.from_pretrained(MODEL_TRANSFORMER, attn_implementation="sdpa", torch_dtype=torch.float16)

# Define ML Model
class FaceEmbeddingModel(torch.nn.Module):
    def __init__(self, model_name, embedding_size):
        super(FaceEmbeddingModel, self).__init__()
        self.config = ViTConfig.from_pretrained(model_name, id2label=idx_to_label, label2id=label_to_idx, return_dict=True)
        self.backbone = ViTModel.from_pretrained(model_name, config=self.config)  # Load ViT model
        self.fc = torch.nn.Linear(self.backbone.config.hidden_size, embedding_size) # Convert to 512D feature vector

    def forward(self, images):
        x = self.backbone(images).last_hidden_state[:, 0]  # Extract embeddings
        x = self.fc(x)  # Convert to 512D embedding
        return torch.nn.functional.normalize(x)  # Normalize for cosine similarity

# Load the model
model_pretrained = torch.load(model_path, map_location=device, weights_only=False)

# Define the ML model - Evaluation function
def prod_function(transformer_model, prod_dl, webcam_dl):
    # Initialize accelerator
    accelerator = Accelerator()

    # to INFO for the main process only.
    if accelerator.is_main_process:
        datasets.utils.logging.set_verbosity_warning()
        transformers.utils.logging.set_verbosity_info()
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

    # The seed need to be set before we instantiate the model, as it will determine the random head.
    set_seed(42)

    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the prepare method.
    accelerated_model, acclerated_prod_dl, acclerated_webcam_dl = accelerator.prepare(transformer_model, prod_dl, webcam_dl)

    # Evaluate at the end of the epoch
    accelerated_model.eval()

    # Find Embedding of the image to be evaluated
    for batch in acclerated_webcam_dl:
        with torch.no_grad():
            #img_prod = acclerated_prod_data['pixel_values']
            emb_prod = accelerated_model(batch['pixel_values'])

    prod_preds = []

    for batch in acclerated_prod_dl:
        #img = batch['pixel_values']
        with torch.no_grad():
            emb = accelerated_model(batch['pixel_values'])
        distance = F.pairwise_distance(emb, emb_prod)

        prod_preds.append(distance)
    return prod_preds

# Creation of Dataloader
class CustomDatasetProd(Dataset):
    def __init__(self, pixel_values):
        self.pixel_values = pixel_values

    def __len__(self):
        return len(self.pixel_values)

    def __getitem__(self, idx):
        item = {
            'pixel_values': self.pixel_values[idx].squeeze(0),
        }
        return item
        
# Creation of Dataset       
class CreateDatasetProd():
    def __init__(self, image_processor):
        super().__init__()
        self.image_processor = image_processor
        # Define a transformation pipeline
        self.transform_prod = transforms.v2.Compose([
                                    transforms.v2.ToImage(),
                                    transforms.v2.ToDtype(torch.uint8, scale=False)
                                ])

    def get_pixels(self, img_paths):
        pixel_values = []
        for path in img_paths:
            # Read and process Images
            img = Image.open(path)
            img = self.transform_prod(img)

            # Scaling the video to ML model's desired format
            img = self.image_processor(img, return_tensors='pt') #, input_data_format='channels_first')

            pixel_values.append(img['pixel_values'].squeeze(0))

            # Force garbage collection
            del img
            gc.collect()
        return pixel_values
        
    def get_pixel(self, img_path):
        # Read and process Images
        img = Image.open(img_path)
        img = self.transform_prod(img)

        # Scaling the video to ML model's desired format
        img = self.image_processor(img, return_tensors='pt') #, input_data_format='channels_first')

        pixel_values = img['pixel_values'] #.squeeze(0)

        # Force garbage collection
        del img
        gc.collect()
        
        return pixel_values
        
    def create_dataset(self, image_paths, webcam=False):
        if webcam == True:
            pixel_values = self.get_pixel(image_paths)
        else:
            pixel_values = torch.stack(self.get_pixels(image_paths))
        
        return CustomDatasetProd(pixel_values=pixel_values)
        
# Read images from directory
image_paths = []
image_file = glob(os.path.join(data_path, '*.jpg'))
#st.write(image_file)
image_paths.extend(image_file)
#st.write('input path size:', len(image_paths))
#st.write(image_paths)

# Create DataLoader for Employees image
dataset_prod_obj = CreateDatasetProd(image_processor_prod)
prod_ds = dataset_prod_obj.create_dataset(image_paths, webcam=False)
prod_dl = DataLoader(prod_ds, batch_size=BATCH_SIZE)

## Testing the dataloader
#prod_inputs = next(iter(prod_dl))
#st.write(prod_inputs['pixel_values'].shape) 

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("# AI Face Recognition app for automated employee attendance")
    
    # About the app Tab
    with gr.Tab("About the app"):
        gr.Markdown(
            """
            ## Product Description/Objective
            An AI face recognition app for automated employee attendance uses advanced facial recognition technology to accurately and efficiently track employee attendance. 
            By simply scanning employees' faces upon arrival and departure, the app eliminates the need for traditional timecards or biometric devices, reducing errors and fraud. 
            It provides real-time attendance data, enhances workplace security, and streamlines HR processes for greater productivity and accuracy.       
            
            ## How does it work ?
            Our app leverages Google's advanced **Vision Transformer (ViT)** architecture, trained on the **LFW (Labeled Faces in the Wild) dataset**, to deliver highly accurate employee attendance tracking through facial recognition. 
            The AI model intelligently extracts distinct facial features and compares them to the stored data of registered employees. When an employee’s face is scanned, the model analyzes the key features, and a confidence score is generated. 
            A high score indicates a match, confirming the employee’s identity and marking their attendance automatically. This seamless, secure process ensures precise tracking while minimizing errors and enhancing workplace efficiency. 
            
            ## About the architecture.

            ## About the Dataset the app is trained on
            """)

    # Gesture recognition Tab
    with gr.Tab("About the app"):
        # Read image from Camera
        enable = st.checkbox("Enable camera")
        picture = st.camera_input("Take a picture", disabled=not enable)
        if picture is not None:
            #img = Image.open(picture)
            #picture.save(webcam_path, "JPEG")
            #st.write('Image saved as:',webcam_path)
        
            ## Create DataLoader for Webcam Image
            webcam_ds = dataset_prod_obj.create_dataset(picture, webcam=True)
            webcam_dl = DataLoader(webcam_ds, batch_size=BATCH_SIZE)
        
            ## Testing the dataloader
            #prod_inputs = next(iter(webcam_dl))
            #st.write(prod_inputs['pixel_values'].shape)
        
            with st.spinner("Wait for it...", show_time=True):
                # Run the predictions
                prediction = prod_function(model_pretrained, prod_dl, webcam_dl)
                predictions = torch.cat(prediction, 0).to(device)
                match_idx = torch.argmin(predictions)
                st.write(predictions)
                st.write(image_paths)
            
                # Display the results
                if predictions[match_idx] <= 0.3:
                  st.write('Welcome: ',image_paths[match_idx].split('/')[-1].split('.')[0])
                else:
                  st.write("Match not found")