svsaurav95 commited on
Commit
78f86c4
·
verified ·
1 Parent(s): ceb6022

Update src/streamlit_app.py

Browse files
Files changed (1) hide show
  1. src/streamlit_app.py +10 -2
src/streamlit_app.py CHANGED
@@ -36,9 +36,17 @@ class MobileViTSegmentation(nn.Module):
36
  # ========== Load Model ==========
37
  @st.cache_resource
38
  def load_model():
39
- checkpoint_path = hf_hub_download(repo_id="svsaurav95/ToothSegmentation", filename="mobilevit_teeth_segmentation.pth")
 
 
 
 
 
 
 
 
40
  model = MobileViTSegmentation(pretrained=False)
41
- model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
42
  model.eval()
43
  return model
44
 
 
36
  # ========== Load Model ==========
37
  @st.cache_resource
38
  def load_model():
39
+ from model import MobileViTSegmentation # or define it here if needed
40
+ cache_dir = "/tmp/huggingface" # Writable path in HF Spaces
41
+
42
+ checkpoint_path = hf_hub_download(
43
+ repo_id="svsaurav95/ToothSegmentation",
44
+ filename="mobilevit_teeth_segmentation.pth",
45
+ cache_dir=cache_dir # <- this fixes the permission issue
46
+ )
47
+
48
  model = MobileViTSegmentation(pretrained=False)
49
+ model.load_state_dict(torch.load(checkpoint_path, map_location="cpu"))
50
  model.eval()
51
  return model
52