Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,6 +7,12 @@ forums](https://discuss.streamlit.io).
 
 In the meantime, below is an example of what you can do with just a few lines of code:
 """
+import os
+
+# Fix for Hugging Face permission error and torch watcher bug
+os.environ["STREAMLIT_HOME"] = "./.streamlit"
+os.environ["STREAMLIT_WATCH_DISABLE"] = "true"
+os.makedirs("./.streamlit", exist_ok=True)
 
 import streamlit as st
 import torch
@@ -24,9 +30,11 @@ from lime.lime_image import LimeImageExplainer
 from skimage.segmentation import mark_boundaries
 import shap
 from shap import GradientExplainer
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 num_classes = 4
 image_size = (224, 224)
+
 # Define CNN Model
 class MyModel(nn.Module):
     def __init__(self, num_classes=4):
@@ -36,27 +44,22 @@ class MyModel(nn.Module):
             nn.BatchNorm2d(64),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=2, stride=2),
-
             nn.Conv2d(64, 128, kernel_size=3, padding=1),
             nn.BatchNorm2d(128),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=2, stride=2),
-
             nn.Conv2d(128, 128, kernel_size=3, padding=1),
             nn.BatchNorm2d(128),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=2, stride=2),
-
             nn.Conv2d(128, 256, kernel_size=3, padding=1),
             nn.BatchNorm2d(256),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=2, stride=2),
-
             nn.Conv2d(256, 256, kernel_size=3, padding=1),
             nn.BatchNorm2d(256),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=2, stride=2),
-
             nn.Conv2d(256, 512, kernel_size=3, padding=1),
             nn.BatchNorm2d(512),
             nn.ReLU(inplace=True),
@@ -67,21 +70,26 @@ class MyModel(nn.Module):
             nn.Linear(512 * 3 * 3, 1024),
             nn.ReLU(inplace=True),
             nn.Dropout(0.25),
-
             nn.Linear(1024, 512),
             nn.ReLU(inplace=True),
             nn.Dropout(0.25),
-
             nn.Linear(512, num_classes)
         )
     def forward(self, x):
         x = self.features(x)
         x = self.classifier(x)
         return x
+
 # Load model
 model = MyModel(num_classes=num_classes).to(device)
-model.load_state_dict(torch.load("src/brainCNNpytorch_model", map_location=torch.device('cpu')))
+try:
+    model.load_state_dict(torch.load("brainCNNpytorch_model", map_location=torch.device('cpu')))
+except FileNotFoundError:
+    st.error("Model file 'brainCNNpytorch_model' not found. Please upload the file correctly.")
+    st.stop()
+
 model.eval()
+
 # Label dictionary
 label_dict = {0: "Meningioma", 1: "Glioma", 2: "No Tumor", 3: "Pituitary"}
 # Preprocessing