updated app

app.py

@@ -2,11 +2,15 @@ import torch
 import gradio as gr
 import matplotlib.pyplot as plt
 from PIL import Image
-from transformers import SamModel, SamProcessor
+from transformers import SamModel, SamConfig, SamProcessor
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# model = SamModel.from_pretrained('hmdliu/sidewalks-seg-base')
+model_config = SamConfig.from_pretrained("facebook/sam-vit-base")
 processor = SamProcessor.from_pretrained('facebook/sam-vit-base')
-model = SamModel.from_pretrained('hmdliu/sidewalks-seg-base')
+model = SamModel(config=model_config)
+model.load_state_dict(torch.load("/content/drive/MyDrive/Project/sidewalk_model_epoch10.pth"))
 model.to(device)
 
 def segment_sidewalk(image, threshold):