Upload app.py

app.py

@@ -1,41 +1,41 @@
+import onnxruntime
+from torchvision import transforms
 import torch
 import torch.nn.functional as F
 import gradio as gr
-import spaces
-import onnxruntime
-from torchvision import transforms
 
-# Enable GPU for ONNX Runtime
 sess_options = onnxruntime.SessionOptions()
 sess_options.enable_profiling = True
-sess_options.add_session_config_entry('session.load_model_format', 'ONNX')
-providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
-ort_sess = onnxruntime.InferenceSession("RFNet.onnx", sess_options=sess_options, providers=providers)
+sess_options.add_session_config_entry('session.load_model_format', 'ONNX') 
+ort_sess = onnxruntime.InferenceSession("RFNet.onnx", sess_options=sess_options)
+
 
 preprocess_img = transforms.Compose([
-    transforms.Resize((352, 352)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-])
+        transforms.Resize((352,352)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])
 
 preprocess_depth = transforms.Compose([
-    transforms.Resize((352, 352)),
-    transforms.ToTensor()
-])
-
-def inference(img, depth, GT):
-    h, w = img.size
-    img = preprocess_img(img).unsqueeze(0).cuda()
-    depth = preprocess_depth(depth.convert('L')).unsqueeze(0).cuda()
-    ort_inputs = {ort_sess.get_inputs()[0].name: img.cpu().numpy(), ort_sess.get_inputs()[1].name: depth.cpu().numpy()}
+        transforms.Resize((352,352)),
+        transforms.ToTensor()])
+def inference(img,depth,GT):
+    h,w = img.size
+    img = preprocess_img(img).unsqueeze(0)
+    depth = preprocess_depth(depth.convert('L')).unsqueeze(0)
+    ort_inputs = {ort_sess.get_inputs()[0].name: img.numpy(), ort_sess.get_inputs()[1].name: depth.numpy()}
     ort_outs = ort_sess.run(None, ort_inputs)
-    output_image = torch.tensor(ort_outs[0]).cuda()
-    res = F.interpolate(output_image, size=(w, h), mode='bilinear', align_corners=False)
-    res = torch.sigmoid(res)
-    res = res.data.cpu().numpy().squeeze()
-    res = (res - res.min()) / (res.max() - res.min() + 1e-8)
+    output_image   = torch.tensor(ort_outs[0])
+    res     = F.interpolate(output_image, size=(w,h), mode='bilinear', align_corners=False)
+    res     = torch.sigmoid(res)
+    res     = res.data.cpu().numpy().squeeze()
+    res     = (res - res.min()) / (res.max() - res.min() + 1e-8)
     return res
 
+
+
+
+
+
 title = "Robust RGB-D Fusion for Saliency Detection"
 description = """ Deployment of the paper: 
 [Robust RGB-D Fusion for Saliency Detection](https://arxiv.org/pdf/2208.01762.pdf) 
@@ -72,4 +72,5 @@ gr.Interface(inference, inputs=[input_1,input_2,input_3], outputs=outputs,
                 description=description,
                 article=article,
                 theme=gr.themes.Soft(),
-                cache_examples=False).launch()
+                cache_examples=False).launch(server_name="0.0.0.0")
+