Update app.py

app.py

@@ -15,6 +15,7 @@ import torch
 
 sns.set_theme()
 
+dir = Path(__file__).resolve().parent
 www_dir = Path(__file__).parent.resolve() / "www"
 
 df = pd.read_csv(Path(__file__).parent / "penguins.csv", na_values="NA")
@@ -64,8 +65,10 @@ def server(input: Inputs, output: Outputs, session: Session):
     model_config = SamConfig.from_pretrained("facebook/sam-vit-base")
     processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
     
-    # Create an instance of the model from my fine-tuned model with the loaded configuration
-    model = SamModel.from_pretrained("aagoluoglu/SAM_Sidewalks", config=model_config)
+    # Create an instance of the model architecture with the loaded configuration
+    model = SamModel(config=model_config)
+    # Update the model by loading the weights from saved file.
+    model.load_state_dict(torch.load(str(dir / "checkpoint.pth")))
     
     # set the device to cuda if available, otherwise use cpu
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -84,7 +87,6 @@ def server(input: Inputs, output: Outputs, session: Session):
         """Displays the uploaded image"""
         img_src = uploaded_image_path()
         if img_src:
-            dir = Path(__file__).resolve().parent
             img: ImgData = {"src": str(dir / uploaded_image_path()), "width": "100px"}
             return img
         else: