Update src/streamlit_app.py

src/streamlit_app.py

@@ -15,6 +15,8 @@ st.markdown("""
     </style>
 """, unsafe_allow_html=True)
 
+MODEL_PATH = mobilevit_teeth_segmentation.pth
+
 # === Model Definition ===
 class MobileViTSegmentation(nn.Module):
     def __init__(self, encoder_name='mobilevit_s', pretrained=False):
@@ -43,7 +45,7 @@ class MobileViTSegmentation(nn.Module):
 @st.cache_resource
 def load_model():
     model = MobileViTSegmentation()
-    state_dict = torch.load("ToothSegmentation/mobilevit_teeth_segmentation.pth", map_location="cpu")
+    state_dict = torch.load( MODEL_PATH, map_location="cpu")
     model.load_state_dict(state_dict)
     model.eval()
     return model