Spaces:
Sleeping
Sleeping
phase1 testing - employees Dataloader
Browse files
app.py
CHANGED
@@ -31,29 +31,152 @@ from glob import glob
|
|
31 |
import shutil
|
32 |
import torch.nn.functional as F
|
33 |
|
|
|
|
|
|
|
34 |
# Initialse Globle Variables
|
35 |
MODEL_TRANSFORMER = 'google/vit-base-patch16-224'
|
36 |
BATCH_SIZE = 8
|
37 |
|
38 |
-
# Set
|
39 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
40 |
|
41 |
-
|
|
|
|
|
|
|
42 |
|
43 |
-
|
44 |
|
|
|
|
|
|
|
45 |
# Read images from directory
|
|
|
|
|
|
|
|
|
|
|
|
|
46 |
|
|
|
|
|
|
|
|
|
47 |
|
48 |
-
#
|
49 |
-
|
50 |
-
|
51 |
-
if picture:
|
52 |
-
col1, col2 = st.columns(2)
|
53 |
-
image = PIL.Image.open(picture)
|
54 |
-
col1.image(image, use_column_width=True)
|
55 |
-
predictions = pipeline(image)
|
56 |
|
57 |
-
|
58 |
-
|
59 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
31 |
import shutil
|
32 |
import torch.nn.functional as F
|
33 |
|
34 |
+
# Set the device (GPU or CPU)
|
35 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
36 |
+
|
37 |
# Initialse Globle Variables
|
38 |
MODEL_TRANSFORMER = 'google/vit-base-patch16-224'
|
39 |
BATCH_SIZE = 8
|
40 |
|
41 |
+
# Set Paths
|
42 |
+
data_path = 'employees'
|
43 |
+
model_path = 'vit_pytorch_GPU_1.pt'
|
44 |
+
|
45 |
+
# Set Title
|
46 |
+
st.title("Employee Attendance System")
|
47 |
+
#pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
|
48 |
+
|
49 |
+
# Define Image Processor
|
50 |
+
image_processor_prod = ViTImageProcessor.from_pretrained(MODEL_TRANSFORMER, attn_implementation="sdpa", torch_dtype=torch.float16)
|
51 |
+
|
52 |
+
# Define ML Model
|
53 |
+
class FaceEmbeddingModel(torch.nn.Module):
|
54 |
+
def __init__(self, model_name, embedding_size):
|
55 |
+
super(FaceEmbeddingModel, self).__init__()
|
56 |
+
self.config = ViTConfig.from_pretrained(model_name, id2label=idx_to_label, label2id=label_to_idx, return_dict=True)
|
57 |
+
self.backbone = ViTModel.from_pretrained(model_name, config=self.config) # Load ViT model
|
58 |
+
self.fc = torch.nn.Linear(self.backbone.config.hidden_size, embedding_size) # Convert to 512D feature vector
|
59 |
+
|
60 |
+
def forward(self, images):
|
61 |
+
x = self.backbone(images).last_hidden_state[:, 0] # Extract embeddings
|
62 |
+
x = self.fc(x) # Convert to 512D embedding
|
63 |
+
return torch.nn.functional.normalize(x) # Normalize for cosine similarity
|
64 |
+
|
65 |
+
# Load the model
|
66 |
+
model_pretrained = torch.load(model_path, map_location=device)
|
67 |
+
|
68 |
+
# Define the ML model - Evaluation function
|
69 |
+
def prod_function(transformer_model, prod_dl, prod_data):
|
70 |
+
# Initialize accelerator
|
71 |
+
accelerator = Accelerator()
|
72 |
+
|
73 |
+
# to INFO for the main process only.
|
74 |
+
if accelerator.is_main_process:
|
75 |
+
datasets.utils.logging.set_verbosity_warning()
|
76 |
+
transformers.utils.logging.set_verbosity_info()
|
77 |
+
else:
|
78 |
+
datasets.utils.logging.set_verbosity_error()
|
79 |
+
transformers.utils.logging.set_verbosity_error()
|
80 |
+
|
81 |
+
# The seed need to be set before we instantiate the model, as it will determine the random head.
|
82 |
+
set_seed(42)
|
83 |
+
|
84 |
+
# 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.
|
85 |
+
accelerated_model, acclerated_prod_dl, acclerated_prod_data = accelerator.prepare(transformer_model, prod_dl, prod_data)
|
86 |
+
|
87 |
+
# Evaluate at the end of the epoch
|
88 |
+
accelerated_model.eval()
|
89 |
+
|
90 |
+
# Find Embedding of the image to be evaluated
|
91 |
+
emb_prod = accelerated_model(acclerated_prod_data)
|
92 |
+
|
93 |
+
prod_preds = []
|
94 |
+
|
95 |
+
for batch in tqdm(acclerated_prod_dl):
|
96 |
+
with torch.no_grad():
|
97 |
+
emb = accelerated_model(**batch)
|
98 |
+
distance = F.pairwise_distance(emb, emb_prod)
|
99 |
+
|
100 |
+
prod_preds.append(distance)
|
101 |
+
return prod_preds
|
102 |
+
|
103 |
+
# Creation of Dataloader
|
104 |
+
class CustomDatasetProd(Dataset):
|
105 |
+
def __init__(self, pixel_values):
|
106 |
+
self.pixel_values = pixel_values
|
107 |
+
|
108 |
+
def __len__(self):
|
109 |
+
return len(self.pixel_values)
|
110 |
+
|
111 |
+
def __getitem__(self, idx):
|
112 |
+
item = {
|
113 |
+
'pixel_values': self.pixel_values[idx].squeeze(0),
|
114 |
+
}
|
115 |
+
return item
|
116 |
+
|
117 |
+
# Creation of Dataset
|
118 |
+
class CreateDatasetProd():
|
119 |
+
def __init__(self, image_processor):
|
120 |
+
super().__init__()
|
121 |
+
self.image_processor = image_processor
|
122 |
+
# Define a transformation pipeline
|
123 |
+
self.transform_prod = transforms.v2.Compose([
|
124 |
+
transforms.v2.ToImage(),
|
125 |
+
transforms.v2.ToDtype(torch.uint8, scale=False)
|
126 |
+
])
|
127 |
+
|
128 |
+
def get_pixels(self, img_paths):
|
129 |
+
pixel_values = []
|
130 |
+
for path in tqdm(img_paths):
|
131 |
+
# Read and process Images
|
132 |
+
img = PIL.Image.open(path)
|
133 |
+
img = self.transform_prod(img)
|
134 |
+
|
135 |
+
# Scaling the video to ML model's desired format
|
136 |
+
img = self.image_processor(img, return_tensors='pt') #, input_data_format='channels_first')
|
137 |
+
|
138 |
+
pixel_values.append(img['pixel_values'].squeeze(0))
|
139 |
|
140 |
+
# Force garbage collection
|
141 |
+
del img
|
142 |
+
gc.collect()
|
143 |
+
return pixel_values
|
144 |
|
145 |
+
def create_dataset(self, image_paths):
|
146 |
|
147 |
+
pixel_values = torch.stack(self.get_pixels(image_paths))
|
148 |
+
return CustomDatasetProd(pixel_values=pixel_values)
|
149 |
+
|
150 |
# Read images from directory
|
151 |
+
image_paths = []
|
152 |
+
image_file = glob(os.path.join(data_path, '*.jpg'))
|
153 |
+
#st.write(image_file)
|
154 |
+
image_paths.extend(image_file)
|
155 |
+
st.write('input path size:', len(image_paths))
|
156 |
+
st.write(image_paths)
|
157 |
|
158 |
+
# Create DataLoader for Employees image
|
159 |
+
dataset_prod_obj = CreateDatasetProd(image_processor_prod)
|
160 |
+
prod_ds = dataset_prod_obj.create_dataset(image_paths)
|
161 |
+
prod_dl = DataLoader(prod_ds, batch_size=BATCH_SIZE)
|
162 |
|
163 |
+
# Testing the dataloader
|
164 |
+
prod_inputs = next(iter(prod_dl))
|
165 |
+
print(prod_inputs['pixel_values'].shape)
|
|
|
|
|
|
|
|
|
|
|
166 |
|
167 |
+
## Read image from Camera
|
168 |
+
#enable = st.checkbox("Enable camera")
|
169 |
+
#picture_path = st.camera_input("Take a picture", disabled=not enable)
|
170 |
+
#if picture_path:
|
171 |
+
# # Create DataLoader for Webcam Image
|
172 |
+
# webcam_ds = dataset_prod_obj.create_dataset(picture_path)
|
173 |
+
# webcam_dl = DataLoader(webcam_ds, batch_size=BATCH_SIZE)
|
174 |
+
#
|
175 |
+
#
|
176 |
+
#prediction = prod_function(model_pretrained, prod_dl, webcam_dl)
|
177 |
+
#predictions = torch.cat(prediction, 0).to('cpu')
|
178 |
+
#match_idx = torch.argmin(predictions)
|
179 |
+
#if predictions[match_idx] <= 0.3:
|
180 |
+
# st.write('Welcome: ',image_paths[match_idx].split('/')[-1].split('.')[0])
|
181 |
+
#else:
|
182 |
+
# st.write("Match not found")
|