Update app.py

app.py

@@ -8,18 +8,59 @@ os.system("pip3 install torch")
 os.system("pip3 install collections")
 os.system("pip3 install torchvision")
 os.system("pip3 install einops")
+os.system("pip3 install pydensecrf")
 #os.system("pip3 install argparse")
-
+import pydensecrf.densecrf as dcrf
 from PIL import Image
 import torch
 from torchvision import transforms
 from model_video import build_model
 import numpy as np
 import collections
+
+def crf_refine(img, annos):
+    def _sigmoid(x):
+        return 1 / (1 + np.exp(-x))
+
+    assert img.dtype == np.uint8
+    assert annos.dtype == np.uint8
+    assert img.shape[:2] == annos.shape
+
+    # img and annos should be np array with data type uint8
+
+    EPSILON = 1e-8
+
+    M = 2  # salient or not
+    tau = 1.05
+    # Setup the CRF model
+    d = dcrf.DenseCRF2D(img.shape[1], img.shape[0], M)
+
+    anno_norm = annos / 255.
+
+    n_energy = -np.log((1.0 - anno_norm + EPSILON)) / (tau * _sigmoid(1 - anno_norm))
+    p_energy = -np.log(anno_norm + EPSILON) / (tau * _sigmoid(anno_norm))
+
+    U = np.zeros((M, img.shape[0] * img.shape[1]), dtype='float32')
+    U[0, :] = n_energy.flatten()
+    U[1, :] = p_energy.flatten()
+
+    d.setUnaryEnergy(U)
+
+    d.addPairwiseGaussian(sxy=3, compat=3)
+    d.addPairwiseBilateral(sxy=60, srgb=5, rgbim=img, compat=5)
+
+    # Do the inference
+    infer = np.array(d.inference(1)).astype('float32')
+    res = infer[1, :]
+
+    res = res * 255
+    res = res.reshape(img.shape[:2])
+    return res.astype('uint8')
+
 #import argparse
 device='cuda:0'
 net = build_model(device).to(device)
-    #net=torch.nn.DataParallel(net)
+#net=torch.nn.DataParallel(net)
 model_path = 'image_best.pth'
 print(model_path)
 weight=torch.load(model_path,map_location=torch.device(device))
@@ -44,8 +85,10 @@ def test(gpu_id, net, img_list, group_size, img_size):
         for i in range(5):
           group_img[i]=img_transform(Image.fromarray(img_list[i]))
         _,pred_mask=net(group_img)
-        print(pred_mask.shape)
-        result = [Image.fromarray((pred_mask[i].detach().squeeze().unsqueeze(2).repeat(1,1,3) * 255).numpy().astype(np.uint8)) for i in range(5)]
+        pred_mask=(pred_mask.detach().squeeze()*255).numpy().astype(np.uint8)
+        pred_mask=[crf_refine(img_list[i],pred_mask[i]) for i in range(5)]
+        print(pred_mask[0].shape)
+        result = [Image.fromarray((torch.from_numpy(pred_mask[i]).unsqueeze(2).repeat(1,1,3)).numpy()) for i in range(5)]
         #w, h = 224,224#Image.open(image_list[i][j]).size
         #result = result.resize((w, h), Image.BILINEAR)
         #result.convert('L').save('0.png')