Update app.py

app.py

@@ -131,9 +131,18 @@ mp_face_mesh = mp.solutions.face_mesh
 face_detection = mp_face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5)
 face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, min_detection_confidence=0.5)
 
-# Load the model from the .pkl file
-model = create_model()  # Ensure architecture matches the saved model
-model.load_state_dict(torch.load("resnet34.pkl", map_location=torch.device('cpu')))  # Use 'cuda' if loading on GPU
+def create_model():
+    model = models.resnet34(pretrained=False)  # Must match the architecture used during training
+    model.fc = torch.nn.Linear(model.fc.in_features, 2)  # Adjust final layer
+    return model
+
+# Load the trained model
+model = create_model()
+try:
+    model.load_state_dict(torch.load("resnet34.pkl", map_location='cpu'))
+except RuntimeError as e:
+    print(f"Error loading model: {e}")
+    # Handle architecture mismatch (e.g., load weights manually)
 model.eval()
 
 # Define transformation for face images