Update app.py

app.py

@@ -147,10 +147,29 @@ class CreateDatasetProd():
             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')
 
-    def create_dataset(self, image_paths):
+        pixel_values = img['pixel_values'].squeeze(0)
 
-        pixel_values = torch.stack(self.get_pixels(image_paths))
+        # Force garbage collection
+        del img
+        gc.collect()
+        
+        return pixel_values
+        
+    def create_dataset(self, image_paths, webcam=False):
+        if webcam == True:
+            pixel_values = torch.stack(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
@@ -163,7 +182,7 @@ image_paths.extend(image_file)
 
 # Create DataLoader for Employees image
 dataset_prod_obj = CreateDatasetProd(image_processor_prod)
-prod_ds = dataset_prod_obj.create_dataset(image_paths)
+prod_ds = dataset_prod_obj.create_dataset(image_paths, webcam=False)
 prod_dl = DataLoader(prod_ds, batch_size=BATCH_SIZE)
 
 # Testing the dataloader
@@ -175,19 +194,19 @@ 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)
+    #picture.save(webcam_path, "JPEG")
+    #st.write('Image saved as:',webcam_path)
 
-    # Create DataLoader for Webcam Image
-    webcam_ds = dataset_prod_obj.create_dataset(webcam_path)
+    ## Create DataLoader for Webcam Image
+    webcam_ds = dataset_prod_obj.create_dataset(webcam_path, webcam=True)
     webcam_dl = DataLoader(webcam_ds, batch_size=BATCH_SIZE)
     
-    # Run the predictions
+    ## Run the predictions
     prediction = prod_function(model_pretrained, prod_dl, webcam_dl)
-    predictions = torch.cat(prediction, 0).to('cpu')
+    predictions = torch.cat(prediction, 0).to(device)
     match_idx = torch.argmin(predictions)
 
-    # Display the results
+    ## Display the results
     if predictions[match_idx] <= 0.3:
       st.write('Welcome: ',image_paths[match_idx].split('/')[-1].split('.')[0])
     else: