Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,6 +7,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torchvision.models as models
+import os
 
 
 def rgb2lab2(r0, g0, b0):
@@ -329,7 +330,8 @@ def prepare_test_image(img, dim=150):
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-model_path = "Hyper_U_NET_pytorch-MAE-30Epoch.pth"
+model_path = os.path.join(os.getcwd(), 'Hyper_U_NET_pytorch-MAE-30Epoch.pth')
+# model_path = "Hyper_U_NET_pytorch-MAE-30Epoch.pth"
 
 test_model = load_model_for_inference(model_path, device)