Update app.py

app.py

@@ -4,8 +4,11 @@ import torchvision
 from torchvision import transforms
 from PIL import Image
 
-
+#####
+###
 # Initialization
+###
+#####
 if 'generate_result' not in st.session_state:
     st.session_state['generate_result'] = 0
 if 'show_result' not in st.session_state:
@@ -13,21 +16,33 @@ if 'show_result' not in st.session_state:
 if 'upload_choice' not in st.session_state:
     st.session_state['upload_choice'] = 'file_up'
 
+
+#####
+###
+# Used to show either the file_uploader or the webcam
+###
+#####
 def change_state():
     if st.session_state['upload_choice'] == 'file_up':
         st.session_state['upload_choice'] = 'webcam'
     else:
         st.session_state['upload_choice'] = 'file_up'
 
-
+# User toggle for file_uploader vs webcam
 st.toggle(label="Webcam",  help="Click on to use webcam, off to upload a file", on_change=change_state)   
 
+# Use state to know whether to show file_uploader or webcam
 if st.session_state['upload_choice'] == 'file_up':
     img = st.file_uploader(label="Upload a photo of a squirrel or bird", type=['png', 'jpg'])
 else:
     st.camera_input(label="Webcam")
-# Load the image and apply transformations
 
+
+#####
+###
+# Load the image and apply transformations
+###
+#####
 def predict_image(image_path, model):
 
     image = Image.open(image_path).convert('RGB')
@@ -61,6 +76,11 @@ def predict_image(image_path, model):
 #         print(f"{class_labels[i]}: {prob:.4f}")
 
 
+#####
+###
+# Load model and prepare for inference
+###
+#####
 model_loaded = torchvision.models.resnet18(pretrained=False)  # Initialize ResNet18 without pretraining 
 model_loaded.fc = torch.nn.Linear(model_loaded.fc.in_features, 2)  # Modify the fully connected layer
 model_loaded = model_loaded.to('cpu')  # Move the model to the appropriate device (GPU or CPU)
@@ -72,6 +92,12 @@ model_loaded.load_state_dict(torch.load(model_path, map_location='cpu'))
 # Set the model to evaluation mode
 model_loaded.eval()
 
+
+#####
+###
+# Toggle view of model output in UI
+###
+#####
 if st.session_state['generate_result'] != 0:
     if img is not None:
         result = predict_image(image_path=img, model=model_loaded)