Update app.py

app.py

@@ -55,6 +55,18 @@ else:
 # Model training and testing in separate directory at ipynb file (Copy of ai-portfolio Kendrick.ipynb)
 
 
+# Load model
+def load_model():
+    model = models.resnet50(weights='DEFAULT')  # Using default weights for initialization
+    num_ftrs = model.fc.in_features
+    model.fc = nn.Linear(num_ftrs, 12)  # Adjust to the number of classes you have
+    
+    # Load the state dict
+    model.load_state_dict(torch.load('resnet50_garbage_classificationv1.2.pth', map_location=torch.device('cpu')))
+    
+    model.eval()  # Set to evaluation mode
+    return model
+
 # Load the model
 model = load_model()