Update app.py

app.py

@@ -165,8 +165,8 @@ prod_ds = dataset_prod_obj.create_dataset(image_paths)
 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) 
+#prod_inputs = next(iter(prod_dl))
+#st.write(prod_inputs['pixel_values'].shape) 
 
 # Read image from Camera
 enable = st.checkbox("Enable camera")
@@ -175,12 +175,14 @@ if picture_path:
     # Create DataLoader for Webcam Image
     webcam_ds = dataset_prod_obj.create_dataset(picture_path)
     webcam_dl = DataLoader(webcam_ds, batch_size=BATCH_SIZE)
-
-
-prediction = prod_function(model_pretrained, prod_dl, webcam_dl)
-predictions = torch.cat(prediction, 0).to('cpu')
-match_idx = torch.argmin(predictions)
-if predictions[match_idx] <= 0.3:
-  st.write('Welcome: ',image_paths[match_idx].split('/')[-1].split('.')[0])
-else:
-  st.write("Match not found")
+    
+    # Run the predictions
+    prediction = prod_function(model_pretrained, prod_dl, webcam_dl)
+    predictions = torch.cat(prediction, 0).to('cpu')
+    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")