Update src/streamlit_app.py

src/streamlit_app.py

@@ -16,13 +16,10 @@ from pathlib import Path
 #########################
 # SETTINGS
 # ########################
-base_path = Path(__file__).resolve().parent.parent
-model_file = base_path / "models" / "momaclassifier_resnet50.pt"
-image_csv = base_path / "data" / "demo_artworks.csv"
-image_folder = base_path / "demo_images"
-model_file = str(model_file)
-image_csv = str(image_csv)
-image_folder = str(image_folder)
+base_path = Path(__file__).resolve().parent
+model_file = str(base_path / "models" / "momaclassifier_resnet50.pt")
+image_csv = str(base_path / "data" / "demo_artworks.csv")
+image_folder = str(base_path / "demo_images")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 class_index = {"Drawing": 0, "Photograph": 1, "Print": 2}
 
@@ -49,7 +46,7 @@ metadata_df = load_metadata()
 @st.cache_resource
 def load_model():
     num_class = len(class_index)
-    model = models.resnet50(pretrained=False)
+    model = models.resnet50(weights=None)
     model.fc = nn.Linear(model.fc.in_features, num_class)
     model.load_state_dict(torch.load(model_file, map_location=device))
     model.to(device)