Update app.py

app.py

@@ -2,16 +2,21 @@ import streamlit as st
 import cv2
 import numpy as np
 import torch
-from torchvision import transforms
+from torchvision import transforms, utils, models
 from PIL import Image
 from TranSalNet_Res import TranSalNet
+from tqdm import tqdm
+import torch.nn as nn
+from utils.data_process import preprocess_img, postprocess_img
+
+
 
 # Load the model and set the device
+device = torch.device('cpu')
 model = TranSalNet()
 model.load_state_dict(torch.load('pretrained_models/TranSalNet_Res.pth', map_location=torch.device('cpu')))
-model.eval()
-device = torch.device('cpu')
 model.to(device)
+model.eval()
 
 # Define Streamlit app
 st.title('Saliency Detection App')