Kaushik066 commited on
Commit
b15f8f4
·
verified ·
1 Parent(s): 7582516

phase1 testing - employees Dataloader

Browse files
Files changed (1) hide show
  1. app.py +138 -15
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 the device (GPU or CPU)
39
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
40
 
41
- pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
 
 
 
42
 
43
- st.title("Hot Dog? Or Not?")
44
 
 
 
 
45
  # Read images from directory
 
 
 
 
 
 
46
 
 
 
 
 
47
 
48
- # Read image from Camera
49
- enable = st.checkbox("Enable camera")
50
- picture = st.camera_input("Take a picture", disabled=not enable)
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
- col2.header("Probabilities")
58
- for p in predictions:
59
- col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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")