# 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) # 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) # Run the predictions prediction = prod_function(model_pretrained, prod_dl, webcam_dl) predictions = torch.cat(prediction, 0).to(device) match_idx = torch.argmin(predictions) # 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")