add sod output

app.py

@@ -9,10 +9,10 @@ import cv2
 # cmd = 'sh download.sh'
 # os.system(cmd)
 
-device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+device = torch.device("cpu")
 print(device)
 inpaint_model = inpaint.IVModel(device=device)
-sod_model = sod.IVModel(device=torch.device("cpu"))
+sod_model = sod.IVModel(device=device)
 max_size=512
 scale_factor = 8
 count = 0
@@ -28,17 +28,19 @@ def sod_inpaint(img):
     w = w // scale_factor * scale_factor
     img = cv2.resize(img, (w,h))
     img = img[:,:,::-1]
-    res = sod_model.forward(img,None)
-    res = np.uint8(res)
-    res = inpaint_model.forward(res,None)
-    res = np.uint8(res)
+    sod_res = sod_model.forward(img,None)
+    sod_res = np.uint8(sod_res)
+    h,w = sod_res.shape[:2]
+    so = np.uint8(sod_res[:,:w//2,:] * (sod_res[:,w//2:,:]/255))
+    inpaint_res = inpaint_model.forward(sod_res,None)
+    inpaint_res = np.uint8(inpaint_res)
     count +=1
     print(count, ' images have been processed')
-    return res[:,:,::-1]
+    return so[:,:,::-1], inpaint_res[:,:,::-1]
 
 
 
 examples = glob.glob('examples/*.*')
 inputs = gr.inputs.Image(shape=None, image_mode="RGB", invert_colors=False, source="upload", tool="editor", type="numpy", label=None, optional=False)
-iface = gr.Interface(fn=sod_inpaint, inputs=inputs, outputs="image", examples=examples, title='Salient Object Detection + Inpaint', description='Upload an image and you will find something disappears', theme='huggingface')
+iface = gr.Interface(fn=sod_inpaint, inputs=inputs, outputs=["image", "image"], examples=examples, title='Salient Object Detection + Inpaint', description='Upload an image and you will see the fg and inpainted bg', theme='huggingface')
 iface.launch()

sod/infer_model.py

@@ -76,7 +76,7 @@ class IVModel():
         img_t = self.input_preprocess_tensor(img)
         shape = [torch.as_tensor([img_t.shape[2]]), torch.as_tensor([img_t.shape[3]])]
         h, w = img_t.shape[2], img_t.shape[3]
-        img_t_temp = F.interpolate(img_t, (512, 512), mode='area')
+        img_t_temp = F.interpolate(img_t, (1024, 1024), mode='area')
         with torch.no_grad():
             res = self.net(img_t_temp, shape=shape)
         res   = F.interpolate(res[0],size=shape, mode='bilinear')